flax

Stop passing reduce_axes to jax.grad, jax.vjp, and jax.value_and_grad.

#4617

Merged