rocm_jax/docs/index.rst

JAX reference documentation
===========================

JAX is Autograd_ and XLA_, brought together for high-performance numerical computing and machine learning research.
It provides composable transformations of Python+NumPy programs: differentiate, vectorize,
parallelize, Just-In-Time compile to GPU/TPU, and more.

.. note::
   JAX 0.4.0 introduces new parallelism APIs, including breaking changes to :func:`jax.experimental.pjit` and a new unified ``jax.Array`` type.
   Please see `Parallelism with JAX <https://jax.readthedocs.io/en/latest/notebooks/Parallelism_with_JAX.html>`_ tutorial and the :ref:`jax-array-migration`
   guide for more information.

.. toctree::
   :maxdepth: 1
   :caption: Getting Started

   installation
   notebooks/quickstart
   notebooks/thinking_in_jax
   notebooks/Common_Gotchas_in_JAX

.. toctree::
   :maxdepth: 1

   jax-101/index

.. toctree::
   :maxdepth: 2

   debugging/index

.. toctree::
   :maxdepth: 1
   :caption: Reference Documentation

   faq
   async_dispatch
   aot
   jaxpr
   notebooks/convolutions
   pytrees
   jax_array_migration
   type_promotion
   errors
   transfer_guard
   glossary
   changelog

.. toctree::
   :maxdepth: 1
   :caption: Advanced JAX Tutorials

   notebooks/autodiff_cookbook
   multi_process
   notebooks/Parallelism_with_JAX
   notebooks/vmapped_log_probs
   notebooks/neural_network_with_tfds_data
   notebooks/Custom_derivative_rules_for_Python_code
   notebooks/How_JAX_primitives_work
   notebooks/Writing_custom_interpreters_in_Jax
   notebooks/Neural_Network_and_Data_Loading
   notebooks/xmap_tutorial
   notebooks/external_callbacks


.. toctree::
   :maxdepth: 1
   :caption: Developer documentation

   contributing
   developer
   jax_internal_api
   autodidax
   jep/index

.. toctree::
   :maxdepth: 1
   :caption: API documentation

   jax

.. toctree::
   :maxdepth: 1
   :caption: Notes

   api_compatibility
   deprecation
   concurrency
   gpu_memory_allocation
   profiling
   device_memory_profiling
   rank_promotion_warning


Indices and tables
==================

* :ref:`genindex`
* :ref:`modindex`
* :ref:`search`


.. _Autograd: https://github.com/hips/autograd
.. _XLA: https://www.tensorflow.org/xla
Add Sphinx-generated reference documentation for JAX. 2019-01-15 20:14:19 -05:00			`JAX reference documentation`
Migrate CHANGELOG.rst -> CHANGELOG.md 2021-02-12 17:03:53 -08:00			`===========================`
Add Sphinx-generated reference documentation for JAX. 2019-01-15 20:14:19 -05:00
DOC: add transformations doc to HTML & reorganize contents 2021-03-08 16:25:04 -08:00			`JAX is Autograd_ and XLA_, brought together for high-performance numerical computing and machine learning research.`
			`It provides composable transformations of Python+NumPy programs: differentiate, vectorize,`
			`parallelize, Just-In-Time compile to GPU/TPU, and more.`
Add Sphinx-generated reference documentation for JAX. 2019-01-15 20:14:19 -05:00
Add jax.Array to the index page 2022-12-01 15:38:38 -08:00			`.. note::`
			JAX 0.4.0 introduces new parallelism APIs, including breaking changes to :func:`jax.experimental.pjit` and a new unified ``jax.Array`` type.
			Please see `Parallelism with JAX <https://jax.readthedocs.io/en/latest/notebooks/Parallelism_with_JAX.html>`_ tutorial and the :ref:`jax-array-migration`
			`guide for more information.`

Notebooks on RTD (#1121) 2019-09-30 11:00:02 -07:00			`.. toctree::`
			`:maxdepth: 1`
DOC: add first JAX-101 notebook 2021-02-26 09:22:34 -08:00			`:caption: Getting Started`
Notebooks on RTD (#1121) 2019-09-30 11:00:02 -07:00
DOC: link to install instructions in HTML docs 2022-04-22 12:35:31 -07:00			`installation`
Notebooks on RTD (#1121) 2019-09-30 11:00:02 -07:00			`notebooks/quickstart`
Add thinking_in_jax.ipynb 2021-01-26 12:08:37 -08:00			`notebooks/thinking_in_jax`
DOC: add first JAX-101 notebook 2021-02-26 09:22:34 -08:00			`notebooks/Common_Gotchas_in_JAX`

			`.. toctree::`
Consolidate links in JAX documentation Move notes down 2022-11-08 14:45:58 -08:00			`:maxdepth: 1`
DOC: add first JAX-101 notebook 2021-02-26 09:22:34 -08:00
			`jax-101/index`
Addressed comments for the Colab. * Cleaned up use of section levels * Renamed ma to multiply_add and sq_add to square_add * Other minor clarifications * Separated the Colabs into Tutorials and Advanced Tutorials 2019-10-03 11:20:04 +02:00
Add debugging tools under `jax.debug` and documentation Co-authored-by: Matthew Johnson <mattjj@google.com> Co-authored-by: Lena Martens <lenamartens@google.com> 2022-07-26 14:47:36 -07:00			`.. toctree::`
			`:maxdepth: 2`

			`debugging/index`

Addressed comments for the Colab. * Cleaned up use of section levels * Renamed ma to multiply_add and sq_add to square_add * Other minor clarifications * Separated the Colabs into Tutorials and Advanced Tutorials 2019-10-03 11:20:04 +02:00			`.. toctree::`
			`:maxdepth: 1`
DOC: add transformations doc to HTML & reorganize contents 2021-03-08 16:25:04 -08:00			`:caption: Reference Documentation`
Addressed comments for the Colab. * Cleaned up use of section levels * Renamed ma to multiply_add and sq_add to square_add * Other minor clarifications * Separated the Colabs into Tutorials and Advanced Tutorials 2019-10-03 11:20:04 +02:00
DOC: add transformations doc to HTML & reorganize contents 2021-03-08 16:25:04 -08:00			`faq`
			`async_dispatch`
write in-process AOT walkthrough doc 2022-09-02 07:46:21 -07:00			`aot`
DOC: add transformations doc to HTML & reorganize contents 2021-03-08 16:25:04 -08:00			`jaxpr`
move convolutions from common gotchas to its own file 2021-02-16 17:21:56 -08:00			`notebooks/convolutions`
DOC: add transformations doc to HTML & reorganize contents 2021-03-08 16:25:04 -08:00			`pytrees`
Add docstrings for `Sharding` classes. Right now I am only documenting `Sharding`, `XLACompatibleSharding`, `MeshPspecSharding` and `SingleDeviceSharding`. Also moving jax_array_migration guide to reference documentation. PiperOrigin-RevId: 488489503 2022-11-14 15:47:06 -08:00			`jax_array_migration`
DOC: add transformations doc to HTML & reorganize contents 2021-03-08 16:25:04 -08:00			`type_promotion`
			`errors`
Make the transfer guard documentation easier to find Move the main documentation of the transfer guard from "Notes" to "Reference documentation" section for better visibility. Add a link to the main documentation to the docstring of jax.transfer_guard(), which currently shows up as the top result when searching for "jax transfer_guard". 2022-06-09 17:56:03 +00:00			`transfer_guard`
DOC: add transformations doc to HTML & reorganize contents 2021-03-08 16:25:04 -08:00			`glossary`
			`changelog`

			`.. toctree::`
			`:maxdepth: 1`
			`:caption: Advanced JAX Tutorials`

DOC: add first JAX-101 notebook 2021-02-26 09:22:34 -08:00			`notebooks/autodiff_cookbook`
Increase visibility in index Multi-Host Multi-Process Envs Guide 2022-11-04 19:22:39 +00:00			`multi_process`
[docs] Rename "JAX in Parallelism" files so the URL matches the title. 2022-12-01 19:53:31 +00:00			`notebooks/Parallelism_with_JAX`
DOC: add first JAX-101 notebook 2021-02-26 09:22:34 -08:00			`notebooks/vmapped_log_probs`
			`notebooks/neural_network_with_tfds_data`
add custom_jvp / vjp, delete custom_transforms 2020-01-15 15:00:38 -08:00			`notebooks/Custom_derivative_rules_for_Python_code`
Added How_JAX_primitives_work colab 2019-10-02 14:41:28 +02:00			`notebooks/How_JAX_primitives_work`
DOC: change build from nbsphinx to myst 2021-02-11 11:56:24 -08:00			`notebooks/Writing_custom_interpreters_in_Jax`
			`notebooks/Neural_Network_and_Data_Loading`
Make xmap tutorial public Note that xmap is still experimental, so treat this more as a preview than as documentation. 2021-03-17 16:12:10 +00:00			`notebooks/xmap_tutorial`
DOC: add example of pure_callback with custom_jvp 2022-12-09 12:43:04 -08:00			`notebooks/external_callbacks`
Notebooks on RTD (#1121) 2019-09-30 11:00:02 -07:00
Add documentation on asynchronous dispatch. 2019-06-04 10:09:43 -04:00
Create developer documentation. * Moved out of README.md some developer-only stuff to docs/developer.rst. * Added documentation about building the documentation 2019-10-03 11:56:57 +02:00			`.. toctree::`
Consolidate links in JAX documentation Move notes down 2022-11-08 14:45:58 -08:00			`:maxdepth: 1`
Create developer documentation. * Moved out of README.md some developer-only stuff to docs/developer.rst. * Added documentation about building the documentation 2019-10-03 11:56:57 +02:00			`:caption: Developer documentation`

Add expanded contributing documentation 2021-04-27 15:00:21 -07:00			`contributing`
Create developer documentation. * Moved out of README.md some developer-only stuff to docs/developer.rst. * Added documentation about building the documentation 2019-10-03 11:56:57 +02:00			`developer`
Added the first draft of the Jaxpr documentation. This replaces the previous Google Doc version, and is now updated with the latest changes in Jaxpr. 2020-02-10 11:40:05 +01:00			`jax_internal_api`
Add autodidax to dev docs toctree and ignore md file to silence warning. Fixes CI failure from google/jax#5827. 2021-02-25 10:46:46 +00:00			`autodidax`
Create JAX Enhancement Proposals (JEPs). Migrate existing design documents to JEPs. 2022-08-08 14:56:08 -04:00			`jep/index`
Create developer documentation. * Moved out of README.md some developer-only stuff to docs/developer.rst. * Added documentation about building the documentation 2019-10-03 11:56:57 +02:00
Add Sphinx-generated reference documentation for JAX. 2019-01-15 20:14:19 -05:00			`.. toctree::`
Consolidate links in JAX documentation Move notes down 2022-11-08 14:45:58 -08:00			`:maxdepth: 1`
Notebooks on RTD (#1121) 2019-09-30 11:00:02 -07:00			`:caption: API documentation`
Add Sphinx-generated reference documentation for JAX. 2019-01-15 20:14:19 -05:00
			`jax`

Consolidate links in JAX documentation Move notes down 2022-11-08 14:45:58 -08:00			`.. toctree::`
			`:maxdepth: 1`
			`:caption: Notes`

			`api_compatibility`
			`deprecation`
			`concurrency`
			`gpu_memory_allocation`
			`profiling`
			`device_memory_profiling`
			`rank_promotion_warning`

Add Sphinx-generated reference documentation for JAX. 2019-01-15 20:14:19 -05:00
			`Indices and tables`
			`==================`

			* :ref:`genindex`
			* :ref:`modindex`
			* :ref:`search`
DOC: add transformations doc to HTML & reorganize contents 2021-03-08 16:25:04 -08:00

			`.. _Autograd: https://github.com/hips/autograd`
Make xmap tutorial public Note that xmap is still experimental, so treat this more as a preview than as documentation. 2021-03-17 16:12:10 +00:00			`.. _XLA: https://www.tensorflow.org/xla`