2022-07-26 14:47:36 -07:00
|
|
|
|
2023-01-18 10:59:42 -08:00
|
|
|
``jax.debug`` module
|
|
|
|
|
====================
|
2022-07-26 14:47:36 -07:00
|
|
|
|
|
|
|
|
.. currentmodule:: jax.debug
|
|
|
|
|
|
|
|
|
|
.. automodule:: jax.debug
|
|
|
|
|
|
2022-09-22 19:04:23 -07:00
|
|
|
Runtime value debugging utilities
|
|
|
|
|
---------------------------------
|
2022-07-26 14:47:36 -07:00
|
|
|
|
2022-09-22 19:04:23 -07:00
|
|
|
:doc:`debugging/print_breakpoint` describes how to make use of JAX's runtime value
|
|
|
|
|
debugging features.
|
2022-07-26 14:47:36 -07:00
|
|
|
|
|
|
|
|
.. autosummary::
|
|
|
|
|
:toctree: _autosummary
|
|
|
|
|
|
|
|
|
|
callback
|
|
|
|
|
print
|
|
|
|
|
breakpoint
|
2022-09-22 19:04:23 -07:00
|
|
|
|
|
|
|
|
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
|
