Debug¶

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

from tinax.debug 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.

visualize_array_sharding ¶

visualize_array_sharding(array: Array) -> None

Print an array's device sharding, rejecting a non-array with a clear TypeError.

array must be concrete, not a value traced under jax.jit: this prints real device placement and cannot run inside a trace.

Parameters:

Name	Type	Description	Default
`array`	`Array`	Concrete JAX array whose sharding is printed.	required

Raises:

Type	Description
`TypeError`	If `array` is not a `jax.Array`.