Home - News ( United States | United Kingdom | Italy | Germany ) - Football scores

Showing content from https://docs.jax.dev/en/latest/_autosummary/jax.profiler.annotate_function.html below:

jax.profiler.annotate_function — JAX documentation

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

Decorator that generates a trace event for the execution of a function.

For example:

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

This will cause an “f” event to show up on the trace timeline if the function execution occurs while the process is being traced by TensorBoard.

Arguments can be passed to the decorator via 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)))

Parameters:

• func (Callable)

• name (str | None)

RetroSearch is an open source project built by @garambo | Open a GitHub Issue

Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo

HTML: 3.2 | Encoding: UTF-8 | Version: 0.7.4