mirror of
https://github.com/ROCm/jax.git
synced 2025-04-19 05:16:06 +00:00
46 lines
785 B
ReStructuredText
46 lines
785 B
ReStructuredText
.. 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
|
|
|
|
Deprecated functions
|
|
--------------------
|
|
|
|
.. autosummary::
|
|
:toctree: _autosummary
|
|
|
|
trace_function
|
|
TraceContext
|
|
StepTraceContext |