jax

Backport approx_tanh fix (cherry-pick jax-ml/jax#34598)

#750

Merged