jax
flax
optax
distrax
rationalnets
softclip
