Add more activation functions from jax.nn.

google / flax

Flax is a neural network library for JAX that is designed for flexibility.

https://flax.readthedocs.io

Apache License 2.0

6.01k stars 633 forks source link

Open carlosgmartin opened 4 months ago

carlosgmartin commented 4 months ago

Adds missing activation functions from jax.nn.