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.

Overloads:
  • fun (None), name (str | None), block_until_ready (bool) → Callable[[Callable[P, T]], Callable[P, T]]

  • fun (Callable[P, T]), name (str | None), block_until_ready (bool) → Callable[P, T]

Parameters:
  • fun (Callable[[~P], T] | None)

  • name (str | None)

  • block_until_ready (bool)

Return type:

Callable[[~P], T] | Callable[[Callable[[~P], T]], Callable[[~P], T]]

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, bound= None)], TypeVar(T)] | None) – 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[[~P], T] | Callable[[Callable[[~P], T]], Callable[[~P], T]]