rocm_jax/docs/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
   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