Content - cbd6584bfd97ebfd5565093ba80e7815bb75025e - fb993e6/docs/jax.profiler.rst

visit type:

Tip revision: 2dfc5f962bfcad645860a4c1fece7280d4ac20bb authored by jax authors on 07 October 2020, 05:59:09 UTC
Merge pull request #4467 from google:update-pypi

Tip revision: 2dfc5f9

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
  trace_function
  TraceContext


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