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

参数: