Fix large constants in activations

jax/nn/functions.py

@@ -35,23 +35,34 @@ def swish(x): return x * sigmoid(x)
 def log_sigmoid(x): return -softplus(-x)
 
 def elu(x, alpha=1.0):
-  safe_x = np.where(x > 0, 0., x)
-  return np.where(x > 0, x, alpha * np.expm1(safe_x))
+  x = np.asarray(x)
+  alpha = np.asarray(alpha, x.dtype)
+  safe_x = lax.select(x > 0, np.zeros(onp.shape(x), x.dtype), x)
+  return lax.select(x > 0, x, alpha * np.expm1(safe_x))
 
 def leaky_relu(x, negative_slope=1e-2):
-  return np.where(x >= 0, x, negative_slope * x)
+  x = np.asarray(x)
+  negative_slope = np.asarray(negative_slope, x.dtype)
+  return lax.select(x >= 0, x, negative_slope * x)
 
 def hard_tanh(x):
-  return np.where(x > 1, 1, np.where(x < -1, -1, x))
+  x = np.asarray(x)
+  shape = onp.shape(x)
+  ones = np.full(shape, 1., x.dtype)
+  minus_ones = np.full(shape, -1., x.dtype)
+  return lax.select(x > 1, ones, lax.select(x < -1, minus_ones, x))
 
 def celu(x, alpha=1.0):
   """Continuously-differentiable exponential linear unit activation"""
-  return np.where(x > 0, x, alpha * np.expm1(x / alpha))
+  x = np.asarray(x)
+  alpha = np.asarray(alpha, x.dtype)
+  return lax.select(x > 0, x, alpha * np.expm1(x / alpha))
 
 def selu(x):
   """Scaled exponential linear unit activation"""
-  alpha = 1.6732632423543772848170429916717
-  scale = 1.0507009873554804934193349852946
+  x = np.asarray(x)
+  alpha = onp.array(1.6732632423543772848170429916717, x.dtype)
+  scale = onp.array(1.0507009873554804934193349852946, x.dtype)
   return scale * elu(x, alpha)
 
 def gelu(x):

tests/nn_test.py

@@ -30,6 +30,7 @@
 from jax import nn
 from jax import random
 import jax
+import jax.numpy as np
 
 from jax.config import config
 config.parse_flags_with_absl()
@@ -50,6 +51,12 @@ def testEluValue(self):
     val = nn.elu(1e4)
     self.assertAllClose(val, 1e4, check_dtypes=False)
 
+  def testEluMemory(self):
+      jax.make_jaxpr(nn.elu)(np.ones((10 ** 12,)))
+
+  def testHardTanhMemory(self):
+      jax.make_jaxpr(nn.hard_tanh)(np.ones((10 ** 12,)))
+
 InitializerRecord = collections.namedtuple(
   "InitializerRecord",
   ["name", "initializer", "shapes"])