neuralqx.profile.decorators module

section(name, cat='', *, args=None, flops=0.0, bytes=0.0)

Context manager to profile a section of code.

Parameters:

  • name (str) – Section name

  • cat (str) – Category (e.g. “vmc”, “sampling”, “jax”, …)

  • args (Optional[Dict[str, Any]]) – Additional metadata to attach to trace events

  • flops (float) – Optional counters for roofline-style derived metrics

  • bytes (float) – Optional counters for roofline-style derived metrics

step(step_num, name='step', cat='step', *, args=None)

Context manager to mark a repeated step (e.g. VMC iteration number).

profile(name=None, cat='', *, sync=None, args=None, flops=0.0, bytes=0.0)

Decorator to profile a function.

sync:

If True, block_until_ready() the returned value inside the timed region. This provides device-time-ish measurements for JAX, but can introduce delays. If None, uses global config (default False).

wrap_callable(fn, *, name=None, cat='', sync=None, args=None, flops=0.0, bytes=0.0, deep=None, deep_cat='python', deep_include=None, deep_exclude=None, deep_max_depth=None)

Wrap an arbitrary callable with profiling sections.

Useful to instrument external library functions without editing their source.

Return type:

Callable[..., Any]

profile_call(fn, *a, name=None, cat='', sync=None, args=None, flops=0.0, bytes=0.0, deep=None, deep_cat='python', deep_include=None, deep_exclude=None, deep_max_depth=None, **k)

Execute one profiled callable invocation.

Return type:

Any

patch_method(obj, method_name, *, name=None, cat='', sync=None, args=None, flops=0.0, bytes=0.0, deep=None, deep_cat='python', deep_include=None, deep_exclude=None, deep_max_depth=None)

Temporarily patch obj.method_name with a profiled wrapper.

patch_attr(target, attr_path, *, name=None, cat='', sync=None, args=None, flops=0.0, bytes=0.0, deep=None, deep_cat='python', deep_include=None, deep_exclude=None, deep_max_depth=None)

Temporarily patch a callable attribute path on an object or module.

attr_path can be dotted, for example "sampler.sample".

python_call_trace(*, cat='python', include_prefixes=None, exclude_prefixes=None, max_depth=None)

Profile nested Python function calls in the current thread.

This is intentionally opt-in and can add noticeable overhead.