Skip to content

← JAX Conceptual Deep-Dives step 4 of 15

Medium Primitives

pure_callback Basics

Jump to the task

Return a host calculation to JAX

Sometimes a needed library function has no suitable JAX implementation. pure_callback lets a pure host function compute a result that compiled JAX code can then use. Before adding a callback, check JAX's own APIs: sigmoid, for example, is already available as jax.nn.sigmoid and jax.scipy.special.expit.

Here we double values with NumPy solely to learn the interface. Production code should simply use 2 * x for this calculation.

import jax
import jax.numpy as jnp
import numpy as np

def host_double(arr):
    return np.asarray(arr) * 2

@jax.jit
def compute(x):
    output = jax.ShapeDtypeStruct(x.shape, x.dtype)
    return jax.pure_callback(host_double, output, x)

result = compute(jnp.array([1., -2., 3.], dtype=jnp.float32))
# result contains [2., -4., 6.]

The callback receives concrete host-side values, not tracers. These can be CPU JAX arrays; np.asarray makes the NumPy boundary explicit. The declaration describes the returned array, not the callback input. Here both have the same shape and dtype. A ShapeDtypeStruct is metadata, not an allocated output buffer, and it does not cast a mismatched result.

Purity is a contract

The same arguments must produce the same result without observable side effects. Do not print, mutate global state or read changing files from a pure callback. JAX can remove an unused call, duplicate it or reorder it when data dependencies allow. Intentional exceptions are also side effects, so do not use this callback for input validation.

A callback is opaque to autodiff: grad does not automatically discover the derivative of the NumPy code. A valid custom JVP or VJP is needed when differentiating through it. Batching also needs an explicit policy, such as vmap_method="sequential" when mapping independent calls. This exercise requires neither gradients nor vmap support.

Crossing the host boundary adds dispatch and possible transfer costs. It does not make the whole program synchronous. Benchmark completed results and compare against a native JAX implementation when available.

Problem

Implement pure_callback_doubler(x):

  1. Define a pure host callback that converts arr with np.asarray and returns the NumPy product arr * 2.
  2. Describe its result with jax.ShapeDtypeStruct(x.shape, x.dtype).
  3. Return jax.pure_callback of that function, declaration and x.

X is a one-dimensional float32 JAX array. Return the same shape and dtype, with doubled values. An empty input returns shape (0,). Return the callback result, not the metadata declaration.

Numerical tests check doubled values; they do not establish that the callback is pure or that its arithmetic used NumPy. Inspect those parts of the implementation too.