Diagnostics¶

tinax.diagnostics sends bounded summaries through injected JAX callbacks and keeps profiler scopes open until result arrays and staged effects complete. Callback execution remains an observational, asynchronous effect.

from tinax.diagnostics import profile_call

result = profile_call("profiles/step-10", train_step, state, batch)

observe_nonfinite ¶

observe_nonfinite(
    values: Array,
    *,
    callback: Callable[[Any, Any], None],
    max_indices: int,
) -> Array

Report a count and bounded row-major flat indices while returning values unchanged.

When non-finite values are present, schedules callback with the total count and up to max_indices row-major flat indices. The callback is an observational, asynchronous effect; values is always returned unchanged.

Parameters:

Name	Type	Description	Default
`values`	`Array`	JAX array to scan for non-finite elements.	required
`callback`	`Callable[[Any, Any], None]`	Host callable invoked as `callback(count, indices)` only when non-finite values exist. Runs asynchronously via `jax.debug.callback`.	required
`max_indices`	`int`	Maximum number of flat indices to report. Must be non-negative; unused slots are filled with `-1`.	required

Returns:

Type	Description
`Array`	The input `values`, unchanged.

Raises:

Type	Description
`TypeError`	If `callback` is not callable, `max_indices` is not an integer (booleans are rejected), or `values` is not a JAX array.
`ValueError`	If `max_indices` is negative.

profile_call ¶

profile_call(
    log_dir: str | PathLike[str],
    function: Callable[P, R],
    /,
    *args: P.args,
    **kwargs: P.kwargs,
) -> R

Profile one call and keep the trace open until result arrays and staged effects complete.

Parameters:

Name	Type	Description	Default
`log_dir`	`str \| PathLike[str]`	Directory the profiler trace is written to.	required
`function`	`Callable[P, R]`	Callable to profile. Invoked once as `function(args, *kwargs)`.	required
`*args`	`P.args`	Positional arguments forwarded to `function`.	`()`
`**kwargs`	`P.kwargs`	Keyword arguments forwarded to `function`.	`{}`

Returns:

Type	Description
`R`	The value returned by `function`, after its result arrays are ready and
`R`	staged effects have completed inside the trace.

Raises:

Type	Description
`TypeError`	If `function` is not callable.