https://github.com/google/jax
Raw File
Tip revision: 9267ac4d75a16ceae7c83eb88edcb1d4891bf843 authored by Michael Levesque-Dion on 21 May 2024, 23:29:43 UTC
Clean up version switches from dense array migration
Tip revision: 9267ac4
advanced_guide.rst
.. _advanced_guide:

Advanced Tutorials
==================
This section contains examples and tutorials on more advanced topics, such as Multi Core computation, Custom operations, and more in depth applications

.. toctree::
   :caption: Examples
   :maxdepth: 1

   notebooks/neural_network_with_tfds_data
   notebooks/Neural_Network_and_Data_Loading
   notebooks/vmapped_log_probs

.. toctree::
   :caption: Parallel Computation
   :maxdepth: 1

   multi_process
   notebooks/Distributed_arrays_and_automatic_parallelization
   notebooks/shard_map
   distributed_data_loading
   notebooks/xmap_tutorial

.. toctree::
   :caption: Automatic Differentiation
   :maxdepth: 1

   notebooks/autodiff_cookbook
   notebooks/Custom_derivative_rules_for_Python_code
   notebooks/autodiff_remat

.. toctree::
   :caption: JAX Internals
   :maxdepth: 1

   notebooks/How_JAX_primitives_work
   notebooks/Writing_custom_interpreters_in_Jax
   Custom_Operation_for_GPUs

.. toctree::
   :caption: Deep Dives
   :maxdepth: 1

   notebooks/convolutions
back to top