JAX reference documentation =============================== Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more For an introduction to JAX, start at the `JAX GitHub page `_. .. toctree:: :maxdepth: 3 :caption: Contents: jax Indices and tables ================== * :ref:`genindex` * :ref:`modindex` * :ref:`search`