- jax.debug.callback(callback, *args, ordered=False, **kwargs)#
Calls a stageable Python callback.
For more explanation, see External Callbacks.
jax.debug.callbackenables you to pass in a Python function that can be called inside of a staged JAX program. A
jax.debug.callbackfollows 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.callbackto 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.
Any) – The positional arguments to the callback.
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.
Any) – The keyword arguments to the callback.
- Return type: