jax.debug.callback#

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

Calls a stageable Python callback.

For more explanation, see External Callbacks.

jax.debug.callback enables you to pass in a Python function that can be called inside of a staged JAX program. A jax.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 jax.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