Add Mish activation function.

PiperOrigin-RevId: 830274208
Change-Id: I463e9042ae3b2c4c7c6ecd3f75111e28144e8d6f

Author: Brax Team

brax/training/networks.py

@@ -47,6 +47,7 @@
     'log_sigmoid': nn.activation.log_sigmoid,
     'log_softmax': nn.activation.log_softmax,
     'logsumexp': nn.activation.logsumexp,
+    'mish': jax.nn.mish,
     'normalize': nn.activation.normalize,
     'one_hot': nn.activation.one_hot,
     'relu': nn.activation.relu,