Getting Started
Reference Documentation
Advanced JAX Tutorials
Notes
custom_vjp
nondiff_argnums
Developer documentation
API documentation
Warning
This page was created from a pull request (#9655).
Methods
__init__(self, arg0, **kwargs)
__init__
is_enabled()
is_enabled
set_metadata(self, **kwargs)
set_metadata
previous
jax.profiler.trace_function
next
jax.profiler.StepTraceContext