jax.profiler.rst
.. currentmodule:: jax.profiler
``jax.profiler`` module
=======================
.. automodule:: jax.profiler
Tracing and time profiling
--------------------------
:doc:`profiling` describes how to make use of JAX's tracing and time profiling
features.
.. autosummary::
:toctree: _autosummary
start_server
start_trace
stop_trace
trace
annotate_function
TraceAnnotation
StepTraceAnnotation
Device memory profiling
-----------------------
See :doc:`device_memory_profiling` for an introduction to JAX's device memory
profiling features.
.. autosummary::
:toctree: _autosummary
device_memory_profile
save_device_memory_profile