Profiling

bartz.profile_mode(value, /)

Context manager to temporarily set profile mode.

Parameters:

value (bool) – Profile mode value to set within the context.

Return type:

Iterator[None]

Examples

>>> with profile_mode(True):
...     # Code runs with profile mode enabled
...     pass

Notes

In profiling mode, the MCMC loop is not compiled into a single function, but instead compiled in smaller pieces that are instrumented to show up in the jax tracer and Python profiling statistics. Search for function names starting with ‘jab’ (see jit_and_block_if_profiling).

Jax tracing is not enabled by this context manager and if used must be handled separately by the user; this context manager only makes sure that the execution flow will be more interpretable in the traces if the tracer is used.