netket.utils.timing.timed

Contents

netket.utils.timing.timed#

netket.utils.timing.timed(fun=None, name=None, block_until_ready=True)[source]#

Marks the decorated function to be timed individually in NetKet timing scopes.

If name is not specified, the qualified name of the function is used.

The profiling is disabled if no global timer is active.

Parameters:

  • fun (Callable[[ParamSpec(P)], TypeVar(T)]) – Function to be decorated

  • name (str | None) – Name to use for the timing of this line.

  • block_until_ready (bool) – Calls jax.block_until_ready on the output if timing (default True), which slows down the code but gives accurate timing results of JAX functions.

Return type:

Callable[[ParamSpec(P)], TypeVar(T)]

Contents