External Callbacks in JAX¶
This guide is a work-in-progress outlining the uses of various callback functions, which allow JAX code to execute certain commands on the host, even while running under jit, vmap, grad, or another transformation.
This is a work-in-progress, and will be updated soon.
TODO(jakevdp, sharadmv): fill-in some simple examples of {func}jax.pure_callback, {func}jax.debug.callback, {func}jax.debug.print, and others.
Example: pure_callback with custom_jvp¶
One powerful way to take advantage of {func}jax.pure_callback is to combine it with {class}jax.custom_jvp (see Custom derivative rules for more details on custom_jvp).
Suppose we want to create a JAX-compatible wrapper for a scipy or numpy function that is not yet available in the jax.scipy or jax.numpy wrappers.
Here, we'll consider creating a wrapper for the Bessel function of the first kind, implemented in scipy.special.jv.
We can start by defining a straightforward pure_callback:
import jax
import jax.numpy as jnp
import scipy.special
def jv(v, z):
v, z = jnp.asarray(v), jnp.asarray(z)
# Require the order v to be integer type: this simplifies
# the JVP rule below.
assert jnp.issubdtype(v.dtype, jnp.integer)
# Promote the input to inexact (float/complex).
# Note that jnp.result_type() accounts for the enable_x64 flag.
z = z.astype(jnp.result_type(float, z.dtype))
# Wrap scipy function to return the expected dtype.
_scipy_jv = lambda v, z: scipy.special.jv(v, z).astype(z.dtype)
# Define the expected shape & dtype of output.
result_shape_dtype = jax.ShapeDtypeStruct(
shape=jnp.broadcast_shapes(v.shape, z.shape),
dtype=z.dtype)
# We use vectorize=True because scipy.special.jv handles broadcasted inputs.
return jax.pure_callback(_scipy_jv, result_shape_dtype, v, z, vectorized=True)
This lets us call into scipy.special.jv from transformed JAX code, including when transformed by jit and vmap:
from functools import partial
j1 = partial(jv, 1)
z = jnp.arange(5.0)
print(j1(z))
WARNING:jax._src.lib.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[ 0. 0.44005057 0.5767248 0.33905897 -0.06604332]
Here is the same result with jit:
print(jax.jit(j1)(z))
[ 0. 0.44005057 0.5767248 0.33905897 -0.06604332]
And here is the same result again with vmap:
print(jax.vmap(j1)(z))
[ 0. 0.44005057 0.5767248 0.33905897 -0.06604332]
However, if we call jax.grad, we see an error because there is no autodiff rule defined for this function:
jax.grad(j1)(z)
--------------------------------------------------------------------------- UnfilteredStackTrace Traceback (most recent call last) <ipython-input-5-fde6421013cd> in <module> ----> 1 jax.grad(j1)(z) /usr/local/lib/python3.8/dist-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs) 161 try: --> 162 return fun(*args, **kwargs) 163 except Exception as e: /usr/local/lib/python3.8/dist-packages/jax/_src/api.py in grad_f(*args, **kwargs) 1090 def grad_f(*args, **kwargs): -> 1091 _, g = value_and_grad_f(*args, **kwargs) 1092 return g /usr/local/lib/python3.8/dist-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs) 161 try: --> 162 return fun(*args, **kwargs) 163 except Exception as e: /usr/local/lib/python3.8/dist-packages/jax/_src/api.py in value_and_grad_f(*args, **kwargs) 1166 if not has_aux: -> 1167 ans, vjp_py = _vjp(f_partial, *dyn_args, reduce_axes=reduce_axes) 1168 else: /usr/local/lib/python3.8/dist-packages/jax/_src/api.py in _vjp(fun, has_aux, reduce_axes, *primals) 2655 flat_fun, out_tree = flatten_fun_nokwargs(fun, in_tree) -> 2656 out_primal, out_vjp = ad.vjp( 2657 flat_fun, primals_flat, reduce_axes=reduce_axes) /usr/local/lib/python3.8/dist-packages/jax/interpreters/ad.py in vjp(traceable, primals, has_aux, reduce_axes) 134 if not has_aux: --> 135 out_primals, pvals, jaxpr, consts = linearize(traceable, *primals) 136 else: /usr/local/lib/python3.8/dist-packages/jax/interpreters/ad.py in linearize(traceable, *primals, **kwargs) 123 jvpfun_flat, out_tree = flatten_fun(jvpfun, in_tree) --> 124 jaxpr, out_pvals, consts = pe.trace_to_jaxpr_nounits(jvpfun_flat, in_pvals) 125 out_primals_pvals, out_tangents_pvals = tree_unflatten(out_tree(), out_pvals) /usr/local/lib/python3.8/dist-packages/jax/_src/profiler.py in wrapper(*args, **kwargs) 313 with TraceAnnotation(name, **decorator_kwargs): --> 314 return func(*args, **kwargs) 315 return wrapper /usr/local/lib/python3.8/dist-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr_nounits(fun, pvals, instantiate) 766 fun = trace_to_subjaxpr_nounits(fun, main, instantiate) --> 767 jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals) 768 assert not env /usr/local/lib/python3.8/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs) 166 try: --> 167 ans = self.f(*args, **dict(self.params, **kwargs)) 168 except: <ipython-input-1-9b5b54cddb29> in jv(v, z) 24 # We use vectorize=True because scipy.special.jv handles broadcasted inputs. ---> 25 return jax.pure_callback(_scipy_jv, result_shape_dtype, v, z, vectorized=True) /usr/local/lib/python3.8/dist-packages/jax/_src/api.py in pure_callback(callback, result_shape_dtypes, *args, **kwargs) 3425 """ -> 3426 return jcb.pure_callback(callback, result_shape_dtypes, *args, **kwargs) 3427 /usr/local/lib/python3.8/dist-packages/jax/_src/callback.py in pure_callback(callback, result_shape_dtypes, vectorized, *args, **kwargs) 130 flat_result_avals, out_tree = tree_util.tree_flatten(result_avals) --> 131 out_flat = pure_callback_p.bind( 132 *flat_args, callback=_flat_callback, /usr/local/lib/python3.8/dist-packages/jax/core.py in bind(self, *args, **params) 328 all(isinstance(arg, Tracer) or valid_jaxtype(arg) for arg in args)), args --> 329 return self.bind_with_trace(find_top_trace(args), args, params) 330 /usr/local/lib/python3.8/dist-packages/jax/core.py in bind_with_trace(self, trace, args, params) 331 def bind_with_trace(self, trace, args, params): --> 332 out = trace.process_primitive(self, map(trace.full_raise, args), params) 333 return map(full_lower, out) if self.multiple_results else full_lower(out) /usr/local/lib/python3.8/dist-packages/jax/interpreters/ad.py in process_primitive(self, primitive, tracers, params) 309 raise NotImplementedError(msg) --> 310 primal_out, tangent_out = jvp(primals_in, tangents_in, **params) 311 if primitive.multiple_results: /usr/local/lib/python3.8/dist-packages/jax/_src/callback.py in pure_callback_jvp_rule(***failed resolving arguments***) 54 del args, kwargs ---> 55 raise ValueError( 56 "Pure callbacks do not support JVP. " UnfilteredStackTrace: ValueError: Pure callbacks do not support JVP. Please use `jax.custom_jvp` to use callbacks while taking gradients. The stack trace below excludes JAX-internal frames. The preceding is the original exception that occurred, unmodified. -------------------- The above exception was the direct cause of the following exception: ValueError Traceback (most recent call last) <ipython-input-5-fde6421013cd> in <module> ----> 1 jax.grad(j1)(z) <ipython-input-1-9b5b54cddb29> in jv(v, z) 23 24 # We use vectorize=True because scipy.special.jv handles broadcasted inputs. ---> 25 return jax.pure_callback(_scipy_jv, result_shape_dtype, v, z, vectorized=True) /usr/local/lib/python3.8/dist-packages/jax/_src/callback.py in pure_callback(callback, result_shape_dtypes, vectorized, *args, **kwargs) 129 lambda x: core.ShapedArray(x.shape, x.dtype), result_shape_dtypes) 130 flat_result_avals, out_tree = tree_util.tree_flatten(result_avals) --> 131 out_flat = pure_callback_p.bind( 132 *flat_args, callback=_flat_callback, 133 result_avals=tuple(flat_result_avals), vectorized=vectorized) /usr/local/lib/python3.8/dist-packages/jax/_src/callback.py in pure_callback_jvp_rule(***failed resolving arguments***) 53 def pure_callback_jvp_rule(*args, **kwargs): 54 del args, kwargs ---> 55 raise ValueError( 56 "Pure callbacks do not support JVP. " 57 "Please use `jax.custom_jvp` to use callbacks while taking gradients.") ValueError: Pure callbacks do not support JVP. Please use `jax.custom_jvp` to use callbacks while taking gradients.
Let's define a custom gradient rule for this. Looking at the definition of the Bessel Function of the First Kind, we find that there is a relatively straightforward recurrence relationship for the derivative with respect to the argument z:
$$ d J_\nu(z) = \left\{ \begin{eqnarray} -J_1(z),\ &\nu=0\\ [J_{\nu - 1}(z) - J_{\nu + 1}(z)]/2,\ &\nu\ne 0 \end{eqnarray}\right. $$
The gradient with respect to $\nu$ is more complicated, but since we've restricted the v argument to integer types we don't need to worry about its gradient for the sake of this example.
We can use jax.custom_jvp to define this automatic differentiation rule for our callback function:
jv = jax.custom_jvp(jv)
@jv.defjvp
def _jv_jvp(primals, tangents):
v, z = primals
_, z_dot = tangents # Note: v_dot is always 0 because v is integer.
jv_minus_1, jv_plus_1 = jv(v - 1, z), jv(v + 1, z)
djv_dz = jnp.where(v == 0, -jv_plus_1, 0.5 * (jv_minus_1 - jv_plus_1))
return jv(v, z), z_dot * djv_dz
Now computing the gradient of our function will work correctly:
j1 = partial(jv, 1)
print(jax.grad(j1)(2.0))
-0.06447162
Further, since we've defined our gradient in terms of jv itself, JAX's architecture means that we get second-order and higher derivatives for free:
jax.hessian(j1)(2.0)
DeviceArray(-0.4003078, dtype=float32, weak_type=True)
Keep in mind that although this all works correctly with JAX, each call to our callback-based jv function will result in passing the input data from the device to the host, and passing the output of scipy.special.jv from the host back to the device.
When running on accelerators like GPU or TPU, this data movement and host synchronization can lead to significant overhead each time jv is called.
However, if you are running JAX on a single CPU (where the "host" and "device" are on the same hardware), JAX will generally do this data transfer in a fast, zero-copy fashion, making this pattern is a relatively straightforward way extend JAX's capabilities.