jax.profiler.annotate_function

jax.profiler.annotate_function#

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

一个装饰器，用于为函数的执行生成跟踪事件。

例如

>>> @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)))

参数: