jax.grad ======== .. currentmodule:: jax .. autofunction:: grad