jax.rst

jax package

Subpackages

.. toctree::
    :maxdepth: 1

    jax.numpy
    jax.scipy
    jax.experimental
    jax.lax
    jax.ops
    jax.random

Module contents

.. automodule:: jax
    :members: jit, disable_jit, grad, value_and_grad, vmap, jacfwd, jacrev, hessian, jvp, linearize, vjp, make_jaxpr
    :undoc-members:
    :show-inheritance: