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.

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

   notebooks/quickstart
   notebooks/thinking_in_jax
   notebooks/Common_Gotchas_in_JAX

.. toctree::
   :maxdepth: 2

   jax-101/index

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

   faq
   async_dispatch
   jaxpr
   notebooks/convolutions
   pytrees
   type_promotion
   errors
   glossary
   changelog

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

   notebooks/autodiff_cookbook
   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
   multi_process

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

   deprecation
   concurrency
   gpu_memory_allocation
   profiling
   device_memory_profiling
   rank_promotion_warning
   custom_vjp_update

.. toctree::
   :maxdepth: 2
   :caption: Developer documentation

   contributing
   developer
   jax_internal_api
   autodidax
   design_notes/index

.. toctree::
   :maxdepth: 3
   :caption: API documentation

   jax


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
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
			`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::`
			`:maxdepth: 2`

			`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
			`.. 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`
			`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`
			`type_promotion`
			`errors`
			`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`
			`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`
Add multi-process guide to docs. This is the first section of the internal doc I drafted. I still need to add the "Common multi-process idioms" and "Multi-host workflows" sections. 2021-05-19 17:22:49 -07:00			`multi_process`
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			`.. toctree::`
			`:maxdepth: 1`
			`:caption: Notes`

Document that JAX follows the NEP-29 deprecation policy. Remove the "experimental" disclaimer from the concurrency documentation. 2021-04-21 11:10:38 -04:00			`deprecation`
Make JAX tracer state thread-local. Allows performing traces in separate threads. Using threading within a traced context still won't work, but that is perhaps less important than the ability to call JIT-ted computations from separate threads. (Revives https://github.com/google/jax/pull/734.) 2019-07-23 09:53:27 -04:00			`concurrency`
Add docs on GPU memory allocation. 2019-07-29 12:24:58 -07:00			`gpu_memory_allocation`
Add some brief documentation about how to profile/trace JAX programs. 2019-08-08 21:02:41 -04:00			`profiling`
Add a heap profiler API and document it. (#3576) 2020-06-26 17:09:09 -04:00			`device_memory_profiling`
improve rank promotion warning, add doc page 2019-08-25 14:28:53 -07:00			`rank_promotion_warning`
Fix some doc build warnings 2021-01-25 10:46:58 -08:00			`custom_vjp_update`
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::`
			`:maxdepth: 2`
			`: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`
Move design notes into docs/, and render them as part of the documentation. 2022-02-02 12:21:48 -05:00			`design_notes/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::`
			`:maxdepth: 3`
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`


			`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`