2020-06-26 17:09:09 -04:00
|
|
|
.. currentmodule:: jax.profiler
|
|
|
|
|
|
2023-01-18 10:59:42 -08:00
|
|
|
``jax.profiler`` module
|
|
|
|
|
=======================
|
2020-03-05 16:07:17 -05:00
|
|
|
|
|
|
|
|
.. automodule:: jax.profiler
|
2020-06-26 17:09:09 -04:00
|
|
|
|
|
|
|
|
Tracing and time profiling
|
|
|
|
|
--------------------------
|
|
|
|
|
|
|
|
|
|
:doc:`profiling` describes how to make use of JAX's tracing and time profiling
|
|
|
|
|
features.
|
|
|
|
|
|
|
|
|
|
.. autosummary::
|
|
|
|
|
:toctree: _autosummary
|
|
|
|
|
|
|
|
|
|
start_server
|
2021-03-17 17:33:25 +00:00
|
|
|
start_trace
|
|
|
|
|
stop_trace
|
|
|
|
|
trace
|
|
|
|
|
annotate_function
|
|
|
|
|
TraceAnnotation
|
|
|
|
|
StepTraceAnnotation
|
2020-06-26 17:09:09 -04:00
|
|
|
|
|
|
|
|
|
|
|
|
|
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
|
