jax package
===========

Subpackages
-----------

.. toctree::
    :maxdepth: 1

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

Just-in-time compilation (:code:`jit`)
--------------------------------------

.. automodule:: jax
    :members: jit, disable_jit, xla_computation, make_jaxpr, eval_shape
    :undoc-members:
    :show-inheritance:

Automatic differentiation
-------------------------

.. automodule:: jax
    :members: grad, value_and_grad, jacfwd, jacrev, hessian, jvp, linearize, vjp, custom_transforms, defjvp, defjvp_all, defvjp, defvjp_all, custom_gradient
    :undoc-members:
    :show-inheritance:


Vectorization (:code:`vmap`)
----------------------------

.. automodule:: jax
    :members: vmap
    :undoc-members:
    :show-inheritance:


Parallelization (:code:`pmap`)
----------------------------

.. automodule:: jax
    :members: pmap
    :undoc-members:
    :show-inheritance:

.. currentmodule:: jax

jax package
===========

Subpackages
-----------

.. toctree::
    :maxdepth: 1

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

Just-in-time compilation (:code:`jit`)
--------------------------------------

.. autofunction:: jit
.. autofunction:: disable_jit
.. autofunction:: xla_computation
.. autofunction:: make_jaxpr
.. autofunction:: eval_shape

Automatic differentiation
-------------------------

.. autofunction:: grad
.. autofunction:: value_and_grad
.. autofunction:: jacfwd
.. autofunction:: jacrev
.. autofunction:: hessian
.. autofunction:: jvp
.. autofunction:: linearize
.. autofunction:: vjp
.. autofunction:: custom_transforms
.. autofunction:: defjvp
.. autofunction:: defjvp_all
.. autofunction:: defvjp
.. autofunction:: defvjp_all
.. autofunction:: custom_gradient


Vectorization (:code:`vmap`)
----------------------------

.. autofunction:: vmap


Parallelization (:code:`pmap`)
------------------------------

.. autofunction:: pmap

Fix index of API docs
Fix index of API docs

The current index of the API docs page seems to have broken links: when I
click on "Automatic differentiation" for example, I get sent to the "JIT"
section.

This change fixes the links.
