Repository: google/jax

Issue:
[Feature Request] mish activation function
Hi, @mattjj @froystig I would like to contribute by adding  ```mish``` activation in  Jax function. It's a  Self Regularized Non-Monotonic Neural Activation Function.

```mish``` =  ```x * tanh(softplus(x))```
paper: https://arxiv.org/abs/1908.08681


Discussion:
This has come up multiple times before, see #1676 and #2842 .

I think the only thing that would change our minds is if `mish` is now more popular. Can you provide evidence that it is widely used? Otherwise it makes sense to me to keep it out of tree, given it is a 1-liner anyway.
Hi, @hawkinsp thanks. ```mish``` is having official implementation in various ML library. Kindly see the attached pics. 
![IMG_20201009_221440](https://user-images.githubusercontent.com/43638955/95610407-e8ed7d00-0a7d-11eb-95fc-3a933e4aef77.jpg)
![IMG_20201009_221425](https://user-images.githubusercontent.com/43638955/95610433-f145b800-0a7d-11eb-9ae2-c11e0d550372.jpg)
This activation function is quite new(2019) but gained a lot of popularity. This paper is having ```82``` citation. And you can find some more amazing stats about ```mish``` here, https://github.com/digantamisra98/Mish .
Thanks.


`mish` seems to have little usage in the research community, so I'm reluctant to add it the standard library. It also wouldn't help JAX users much to add it; if anyone wants `mish`, the complete implementation is:
```
def mish(x): return x * jnp.tanh(jax.nn.softplus(x))
```
and that's something that is easy to copy and paste into a model.

We can reevaluate if `mish` has significant research impact and becomes widely used. So I would say: not yet.
@hawkinsp It would be nice to re-evaluate this request as of now considering Mish got added to standard PyTorch ([PR](https://github.com/pytorch/pytorch/pull/58648), [Feature Request Issue](https://github.com/pytorch/pytorch/issues/58375#issuecomment-842654130)) based on similar evaluation criteria. 
Additionally the paper has received nearly 270 citations as per [Scholar](https://scholar.google.com/citations?user=LwiJwNYAAAAJ&hl=en).
The paper now has 1,663 citations according to [Google Scholar](https://scholar.google.com/scholar?q=mish+activation+function).

Some places where mish has been requested:
- https://github.com/google/flax/issues/584
- https://github.com/google/flax/discussions/3391
- https://github.com/google/objax/issues/225
- https://github.com/google/jax/issues/2842
- https://github.com/google/flax/discussions/3483
- https://github.com/google/jax/pull/1676
- https://github.com/google/jax/issues/4523

Looks like there's enough demand and citations for it.

The following links might be of interest:
- [Benchmarking PyTorch’s Native Mish](https://benjaminwarner.dev/2021/07/19/benchmarking-pytorch-native-mish)
- [mish-cuda: Mish Activation Function for PyTorch](https://github.com/thomasbrandon/mish-cuda)

So there might be value in an optimized implementation.

Resolving patch:
diff --git a/jax/_src/nn/functions.py b/jax/_src/nn/functions.py
--- a/jax/_src/nn/functions.py
+++ b/jax/_src/nn/functions.py
@@ -199,6 +199,29 @@ def silu(x: ArrayLike) -> Array:
 
 swish = silu
 
+@jax.jit
+def mish(x: ArrayLike) -> Array:
+  r"""Mish activation function.
+
+  Computes the element-wise function:
+
+  .. math::
+    \mathrm{mish}(x) = x \cdot \mathrm{tanh}(\mathrm{softplus}(x))
+
+  For more information, see
+  `Mish: A Self Regularized Non-Monotonic Activation Function
+  <https://arxiv.org/abs/1908.08681>`_.
+
+  Args:
+    x : input array
+
+  Returns:
+    An array.
+  """
+  numpy_util.check_arraylike("mish", x)
+  x_arr = jnp.asarray(x)
+  return x_arr * jnp.tanh(softplus(x_arr))
+
 @jax.jit
 def log_sigmoid(x: ArrayLike) -> Array:
   r"""Log-sigmoid activation function.
@@ -314,7 +337,7 @@ def celu(x: ArrayLike, alpha: ArrayLike = 1.0) -> Array:
 
   For more information, see
   `Continuously Differentiable Exponential Linear Units
-  <https://arxiv.org/pdf/1704.07483.pdf>`_.
+  <https://arxiv.org/abs/1704.07483>`_.
 
   Args:
     x : input array
@@ -342,7 +365,7 @@ def selu(x: ArrayLike) -> Array:
 
   For more information, see
   `Self-Normalizing Neural Networks
-  <https://papers.nips.cc/paper/6698-self-normalizing-neural-networks.pdf>`_.
+  <https://arxiv.org/abs/1706.02515>`_.
 
   Args:
     x : input array
diff --git a/jax/nn/__init__.py b/jax/nn/__init__.py
--- a/jax/nn/__init__.py
+++ b/jax/nn/__init__.py
@@ -45,6 +45,7 @@
   silu as silu,
   swish as swish,
   squareplus as squareplus,
+  mish as mish,
 )
 
 # Deprecations


Test patch:
diff --git a/tests/nn_test.py b/tests/nn_test.py
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -91,6 +91,26 @@ def testSquareplusGradNan(self):
   def testSquareplusZero(self, dtype):
     self.assertEqual(dtype(1), nn.squareplus(dtype(0), dtype(4)))
 
+  def testMishGrad(self):
+    check_grads(nn.mish, (1e-8,), order=4,
+                rtol=1e-2 if jtu.test_device_matches(["tpu"]) else None)
+
+  def testMishGradZero(self):
+    check_grads(nn.mish, (0.,), order=1,
+                rtol=1e-2 if jtu.test_device_matches(["tpu"]) else None)
+
+  def testMishGradNegInf(self):
+    check_grads(nn.mish, (-float('inf'),), order=1,
+                rtol=1e-2 if jtu.test_device_matches(["tpu"]) else None)
+
+  def testMishGradNan(self):
+    check_grads(nn.mish, (float('nan'),), order=1,
+                rtol=1e-2 if jtu.test_device_matches(["tpu"]) else None)
+
+  @parameterized.parameters([float] + jtu.dtypes.floating)
+  def testMishZero(self, dtype):
+    self.assertEqual(dtype(0), nn.mish(dtype(0)))
+
   def testReluGrad(self):
     rtol = 1e-2 if jtu.test_device_matches(["tpu"]) else None
     check_grads(nn.relu, (1.,), order=3, rtol=rtol)
@@ -117,6 +137,10 @@ def testSquareplusValue(self):
     val = nn.squareplus(1e3)
     self.assertAllClose(val, 1e3, check_dtypes=False, atol=1e-3)
 
+  def testMishValue(self):
+    val = nn.mish(1e3)
+    self.assertAllClose(val, 1e3, check_dtypes=False, atol=1e-3)
+
   @jtu.skip_on_flag("jax_skip_slow_tests", True)
   def testEluGrad(self):
     check_grads(nn.elu, (1e4,), order=4, eps=1.)
@@ -149,7 +173,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.sparse_plus, nn.sigmoid, nn.squareplus)))
+       nn.relu, nn.softplus, nn.sparse_plus, nn.sigmoid, nn.squareplus, nn.mish)))
   def testDtypeMatchesInput(self, dtype, fn):
     x = jnp.zeros((), dtype=dtype)
     out = fn(x)