jax.profiler.annotate_function

jax.profiler.annotate_function#

jax.profiler.annotate_function(func, name=None, **decorator_kwargs)[source]#

为函数执行生成跟踪事件的装饰器。

例如

>>> @jax.profiler.annotate_function
... def f(x):
...   return jnp.dot(x, x.T).block_until_ready()
>>>
>>> result = f(jnp.ones((1000, 1000)))

如果函数执行发生在进程被 TensorBoard 跟踪时，这将导致在跟踪时间轴上显示“f”事件。

可以通过 functools.partial() 将参数传递给装饰器。

>>> from functools import partial

>>> @partial(jax.profiler.annotate_function, name="event_name")
... def f(x):
...   return jnp.dot(x, x.T).block_until_ready()

>>> result = f(jnp.ones((1000, 1000)))
参数:

  • func (Callable)

  • name (str | None | None)

内容