Add sparseplus activation to jax.nn.

PiperOrigin-RevId: 616087452

Author: mtthss

docs/jax.nn.rst

@@ -22,6 +22,7 @@ Activation functions
     relu6
     sigmoid
     softplus
+    sparse_plus
     soft_sign
     silu
     swish

jax/_src/nn/functions.py

@@ -110,6 +110,31 @@ def softplus(x: ArrayLike) -> Array:
   """
   return jnp.logaddexp(x, 0)
 
+@jax.jit
+def sparse_plus(x: ArrayLike) -> Array:
+  r"""Sparse plus function.
+
+  Computes the function:
+
+  .. math::
+
+    \mathrm{sparse\_plus}(x) = \begin{cases}
+      0, & x \leq -1\\
+      \frac{1}{4}(x+1)^2, & -1 < x < 1 \\
+      x, & 1 \leq x
+    \end{cases}
+
+  This is the twin function of the softplus activation ensuring a zero output
+  for inputs less than -1 and a linear output for inputs greater than 1,
+  while remaining smooth, convex, monotonic by an adequate definition between
+  -1 and 1.
+
+  Args:
+    x: input (float)
+  """
+  x = jnp.asarray(x)
+  return jnp.where(x <= -1.0, 0.0, jnp.where(x >= 1.0, x, (x + 1.0)**2/4))
+
 @jax.jit
 def soft_sign(x: ArrayLike) -> Array:
   r"""Soft-sign activation function.

jax/nn/__init__.py

@@ -41,6 +41,7 @@
   soft_sign as soft_sign,
   softmax as softmax,
   softplus as softplus,
+  sparse_plus as sparse_plus,
   silu as silu,
   swish as swish,
   squareplus as squareplus,

tests/nn_test.py

@@ -63,6 +63,14 @@ def testSoftplusGradNan(self):
   def testSoftplusZero(self, dtype):
     self.assertEqual(jnp.log(dtype(2)), nn.softplus(dtype(0)))
 
+  def testSparseplusGradZero(self):
+    check_grads(nn.sparse_plus, (-2.,), order=1,
+                rtol=1e-2 if jtu.test_device_matches(["tpu"]) else None)
+
+  def testSparseplusGrad(self):
+    check_grads(nn.sparse_plus, (0.,), order=1,
+                rtol=1e-2 if jtu.test_device_matches(["tpu"]) else None)
+
   def testSquareplusGrad(self):
     check_grads(nn.squareplus, (1e-8,), order=4,
                 rtol=1e-2 if jtu.test_device_matches(["tpu"]) else None)
@@ -101,6 +109,10 @@ def testSoftplusValue(self):
     val = nn.softplus(89.)
     self.assertAllClose(val, 89., check_dtypes=False)
 
+  def testSparseplusValue(self):
+    val = nn.sparse_plus(89.)
+    self.assertAllClose(val, 89., check_dtypes=False)
+
   def testSquareplusValue(self):
     val = nn.squareplus(1e3)
     self.assertAllClose(val, 1e3, check_dtypes=False, atol=1e-3)
@@ -137,7 +149,7 @@ def gelu_reference(x):
       (jnp.float32, jnp.bfloat16, jnp.float16),
       (partial(nn.gelu, approximate=False),
        partial(nn.gelu, approximate=True),
-       nn.relu, nn.softplus, nn.sigmoid, nn.squareplus)))
+       nn.relu, nn.softplus, nn.sparse_plus, nn.sigmoid, nn.squareplus)))
   def testDtypeMatchesInput(self, dtype, fn):
     x = jnp.zeros((), dtype=dtype)
     out = fn(x)