jax.debug.callback

jax.debug.callback(callback, *args, ordered=False, **kwargs)

Calls a stageable Python callback.

For more explanation, see External Callbacks.

debug_callback enables you to pass in a Python function that can be called inside of a staged JAX program. A debug_callback follows existing JAX transformation pure operational semantics, which are therefore unaware of side-effects. This means the effect could be dropped, duplicated, or potentially reordered in the presence of higher-order primitives and transformations.

We want this behavior because we’d like debug_callback to be “innocuous”, i.e. we want these primitives to change the JAX computation as little as possible while revealing as much about them as possible, such as which parts of the computation are duplicated or dropped.

Parameters

• callback (Callable[..., Any]) – A Python callable. Its return value will be ignored.

• *args (Any) – The positional arguments to the callback.

• ordered (bool) – A keyword only argument used to indicate whether or not the staged out computation will enforce ordering of this callback w.r.t. other ordered callbacks.

• **kwargs (Any) – The keyword arguments to the callback.

Return type

None

Returns

None

See also

• jax.experimental.io_callback(): callback designed for impure functions.

• jax.pure_callback(): callback designed for pure functions.

• jax.debug.print(): callback designed for printing.