netket.utils.timing.timed_scope

Contents

netket.utils.timing.timed_scope#

netket.utils.timing.timed_scope(name=None, force=False)[source]#

Context manager used to mark a scope to be timed individually by NetKet timers.

If name is not specified, the file name and line number is used.

If force is not specified, the timer only runs if a top-level timer is in use as well. If force is specified, the timer and nested timers will always run.

Note

If you are using JAX functions in your code, to get reliable timing results you should block the output of the function with block_until_ready. As this can sensibly slow down your code, this is only done when force=True or if a top-level timer is in use.

To block the output on those conditions, simply use timer.block_until_ready(output) on the output of the function you want to time.

Example

Time a section of code

>>> import netket as nk
>>> import time; import jax
>>>
>>> with nk.utils.timing.timed_scope(force=True) as timer:
...    time.sleep(1)  # This line and the ones below are indented
...    with nk.utils.timing.timed_scope("subfunction 1"):
...       time.sleep(0.5)
...    with nk.utils.timing.timed_scope("subfunction 2"):
...       time.sleep(0.25)
...    a = jax.random.normal(jax.random.key(1), (100,100))
...    # Must block jax functions otherwise the timing is off
...    timer.block_until_ready(timer)
...
>>>
>>> timer
╭──────────────────────── Timing Information ─────────────────────────╮
│ Total: 1.763                                                        │
│ ├── (28.7%) | subfunction 1 : 0.505 s                               │
│ └── (14.3%) | subfunction 2 : 0.252 s                               │
╰─────────────────────────────────────────────────────────────────────╯
Parameters:

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

  • force (bool) – whether to always time, even if no top level timer is in use

Contents