Content - 8e3c29ba8412f380b7f7ad9241505a7d38f17910 - 83c1049/docs/jax.profiler.rst

visit type:

Tip revision: 4ed8f5a895e799237ff17b9a6574a15f454a1776 authored by Tom Cobley on 11 July 2024, 17:52:02 UTC
Allow implementations of XLA:CPU collectives to be passed directly into `make_cpu_client`.

Tip revision: 4ed8f5a

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

Browse the archive

https://github.com/google/jax