Use scipy.special.expit for sigmoid

tensor2tensor/trax/backend.py

@@ -158,6 +158,7 @@ def jax_randint(key, shape, minval, maxval, dtype=onp.int32):
     "name": "jax",
     "np": jnp,
     "logsumexp": jax_special.logsumexp,
+    "expit": jax_special.expit,
     "conv": jax_conv,
     "avg_pool": jax_avg_pool,
     "max_pool": jax_max_pool,

tensor2tensor/trax/layers/core.py

@@ -35,7 +35,7 @@ def Relu(x, **unused_kwargs):
 
 @base.layer()
 def Sigmoid(x, **unused_kwargs):
-  return 1. / (1. + np.exp(-x))
+  return backend.expit(x)
 
 
 @base.layer()

tensor2tensor/trax/layers/initializers.py

@@ -22,14 +22,9 @@
 from tensor2tensor.trax import backend
 
 
-<<<<<<< HEAD
-def _get_fans(shape, out_dim=-1, in_dim=-2):
-  #temporary fix until numpy.delete supports negative indices
-=======
 def _GetFans(shape, out_dim=-1, in_dim=-2):
   """Get the fan-in and fan-out sizes for the given shape and dims."""
   # Temporary fix until numpy.delete supports negative indices.
->>>>>>> 8edb68ca31a02fe96a591e2ca4bd38a0e447277b
   if out_dim < 0:
     out_dim += len(shape)
   if in_dim < 0:
@@ -39,17 +34,11 @@ def _GetFans(shape, out_dim=-1, in_dim=-2):
   if len(shape) >= 2:
     fan_in, fan_out = shape[in_dim], shape[out_dim]
   elif len(shape) == 1:
-<<<<<<< HEAD
-    fan_in, fan_out = shape[0]
-  else:
-    fan_in, fan_out = 1.
-=======
     fan_in = shape[0]
     fan_out = shape[0]
   else:
     fan_in = 1.
     fan_out = 1.
->>>>>>> 8edb68ca31a02fe96a591e2ca4bd38a0e447277b
     fan_in *= receptive_field
     fan_out *= receptive_field
   return fan_in, fan_out
@@ -75,11 +64,7 @@ def Init(shape, rng):
 
 
 def VarianceScalingInitializer(out_dim, in_dim, scale, mode, distribution):
-<<<<<<< HEAD
-  """Initializer capable of adapting its scale to the shape of weights tensors."""
-=======
   """Initializer capable of adapting its scale to the shape of weights."""
->>>>>>> 8edb68ca31a02fe96a591e2ca4bd38a0e447277b
   if scale <= 0.:
     raise ValueError('scale must be positive float, {} given'.format(scale))
   if mode not in {'fan_in', 'fan_out', 'fan_avg'}:
@@ -88,12 +73,8 @@ def VarianceScalingInitializer(out_dim, in_dim, scale, mode, distribution):
         .format(mode))
 
   def Init(shape, rng):
-<<<<<<< HEAD
-    fan_in, fan_out = _get_fans(shape, out_dim, in_dim)
-=======
     """The initializer function."""
     fan_in, fan_out = _GetFans(shape, out_dim, in_dim)
->>>>>>> 8edb68ca31a02fe96a591e2ca4bd38a0e447277b
     gain = scale
     if mode == 'fan_in':
       gain /= fan_in