jax.profiler.StepTraceContext

class jax.profiler.StepTraceContext(*args, **kwargs)

__init__(self: jaxlib.xla_extension.profiler.TraceMe, arg0: str, **kwargs) → None

Methods

__init__(self, arg0, **kwargs)
is_enabled()
set_metadata(self, **kwargs)