Skip to content

← JAX Conceptual Deep-Dives step 5 of 15

Hard Primitives

io_callback for Side Effects

Jump to the task

Tell JAX when host work has side effects

A pure callback can be removed when its result is unused. That is unsuitable for a required host action such as recording a measurement. io_callback represents impure host work and can also return values to the JAX calculation. Its output shape and dtype must be declared.

The example below records a scalar in a local list. It makes no file or network requests.

import jax
import jax.numpy as jnp
from jax.experimental import io_callback

recorded = []

def record(value):
    recorded.append(float(value))
    return value

@jax.jit
def compute(x):
    total = jnp.sum(x)
    output = jax.ShapeDtypeStruct(total.shape, total.dtype)
    return io_callback(record, output, total, ordered=True)

result = compute(jnp.array([1., 2., 3.]))
result.block_until_ready()
jax.effects_barrier()
# result is 6; recorded is [6.0]

What ordering means

ordered=True preserves order between ordered I/O callbacks. It does not impose a global sequence on unrelated device operations or make an external write transactional. The default is ordered=False; choose ordering when the relative order of the host actions matters.

An I/O callback is retained even if its result is unused, unlike a pure callback. This is a compiler execution property, not a guarantee against process failures or retries in an external system.

Ordered I/O callbacks cannot be directly vmapped. Grad also does not support differentiation through an I/O callback receiving differentiated values. For training diagnostics, consider returning metrics and handling them in the Python loop after the differentiated calculation.

Host work and transfers cost time, but callbacks do not universally turn asynchronous dispatch off. Profile the actual workload; avoid adding a callback to every small operation without a reason.

Problem

Implement io_callback_passthrough(x):

  1. Define a host function that returns arr * 2.
  2. Declare the result with jax.ShapeDtypeStruct(x.shape, x.dtype).
  3. Call io_callback with that function, declaration, x and ordered=True.
  4. Return its result.

X is a one-dimensional float32 JAX array. Return doubled values with the same shape and dtype. [5, 6, 7] becomes [10, 12, 14]; an empty array stays empty. Despite the entrypoint name, this is doubling, not an unchanged pass-through.

The exercise practices I/O-callback syntax without requiring a real side effect. Do not add file writes, logs or external requests to the submission. For a pure doubler outside this exercise, ordinary JAX arithmetic is preferable. Numerical tests cannot verify ordering or prove that an I/O callback was used; inspect the API call as well.