https://github.com/google/jax
Raw File
Tip revision: 0f30685dac17c509afa7e6c4a6a9341011dfe7e4 authored by Skye Wanderman-Milne on 09 August 2023, 17:05:46 UTC
Remove StreamExecutor-based TPU runtime from Cloud TPU CI
Tip revision: 0f30685
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/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