jax.profiler.TraceAnnotation

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

一个在性能分析器中生成跟踪事件的上下文管理器。

此跟踪事件覆盖上下文所包含代码的持续时间。

例如

>>> x = jnp.ones((1000, 1000))
>>> with jax.profiler.TraceAnnotation("my_label"):
...   result = jnp.dot(x, x.T).block_until_ready()

如果该事件在进程被跟踪期间发生，这会使一个名为“my_label”的事件显示在跟踪时间轴上。

__init__(self, arg0: str, /, **kwargs) → None

属性

is_enabled
set_metadata