neuralqx.profile.profiler module

class Profiler

Bases: object

Process-local profiler with hierarchical section timing and optional trace/telemetry.

Thread-safety:

  • Each thread gets an independent stack to compute exclusive time.

  • Summary aggregation merges into a shared tree. Updates are guarded by a lock. (Summary updates happen at section exit only; trace emission is optional.)

enabled()
Return type:

bool

refresh_if_needed()

Refresh config if PROFILE toggled at runtime (common in notebooks).

This is designed to be extremely cheap in the steady state: it only reads os.environ[“NQX_PROFILE”] via profiling_enabled().

Return type:

None

maybe_sync(value)

Synchronize JAX arrays if configured. This can introduce delays.

Return type:

Any

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

Create a profiling section context manager.

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

Mark a repeated step. Adds metadata and (optionally) JAX step trace annotation.

summary_dict()
Return type:

Dict[str, Any]

trace_events()
Return type:

List[Dict[str, Any]]

flush()
Return type:

None

sync_enabled()

Allow runtime toggling of sync via env var without recreating the Profiler.

Return type:

bool

get_profiler()
Return type:

Profiler | _DisabledProfiler