Revision 6c05add8ae69fb25e09b3bb537ea90a2a3033b95 authored by yashkatariya on 15 November 2022, 20:59:00 UTC, committed by yashkatariya on 15 November 2022, 20:59:00 UTC
1 parent 233bf21
Raw File
jax.debug.rst

jax.debug package
=================

.. currentmodule:: jax.debug

.. automodule:: jax.debug

Runtime value debugging utilities
---------------------------------

:doc:`debugging/print_breakpoint` describes how to make use of JAX's runtime value
debugging features.

.. autosummary::
  :toctree: _autosummary

  callback
  print
  breakpoint

Sharding debugging utilities
----------------------------

Functions that enable inspecting and visualizing array shardings inside (and outside)
staged functions.

.. autosummary::
  :toctree: _autosummary

  inspect_array_sharding
  visualize_array_sharding
  visualize_sharding
back to top