jax-ml/jax · error · RuntimeError

jax.pure_callback failed to find a local CPU device to place

Error message

jax.pure_callback failed to find a local CPU device to place the inputs on. Make sure "cpu" is listed in --jax_platforms or the JAX_PLATFORMS environment variable.

What it means

pure_callback_impl executes the callback on a local CPU device via OpSharding (collective broadcast) and therefore needs to find one with xb.local_devices(backend='cpu'). If the CPU backend has no visible local devices (common when JAX_PLATFORMS restricts platforms, or in a process pinned to accelerators only), the lookup raises RuntimeError and this wrapper re-raises with guidance.

Source

Thrown at jax/_src/callback.py:85

  in_tree: tree_util.PyTreeDef  # (args, kwargs) pytree for `callback_func`.

  def __call__(self, *flat_args: Array) -> Sequence[Array]:
    args, kwargs = tree_util.tree_unflatten(self.in_tree, flat_args)
    return tree_util.tree_leaves(self.callback_func(*args, **kwargs))


def pure_callback_impl(
    *args,
    result_avals,
    callback: _FlatCallback,
    sharding: Sharding | None,
    vmap_method: str | None,
):
  del sharding, vmap_method, result_avals
  try:
    cpu_device, *_ = xb.local_devices(backend="cpu")
  except RuntimeError as e:
    raise RuntimeError(
        "jax.pure_callback failed to find a local CPU device to place the"
        " inputs on. Make sure \"cpu\" is listed in --jax_platforms or the"
        " JAX_PLATFORMS environment variable."
    ) from e
  args = api.device_put(args, cpu_device)
  with config.default_device(cpu_device):
    try:
      return tree_util.tree_map(np.asarray, callback(*args))
    except BaseException:
      logger.exception("jax.pure_callback failed")
      raise


pure_callback_p.def_impl(functools.partial(dispatch.apply_primitive,
                                           pure_callback_p))


@pure_callback_p.def_abstract_eval

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Add 'cpu' to the platform list: JAX_PLATFORMS=cpu,cuda or jax.config.update('jax_platforms', ('cpu','cuda'))
  2. Unset JAX_PLATFORMS entirely if CPU restriction is not required
  3. If running pure_callback logic on GPU is acceptable, restructure to move data back to host with jax.device_get outside jit instead of a callback

Example fix

# before
JAX_PLATFORMS=cuda python train.py  # pure_callback -> RuntimeError

# after
JAX_PLATFORMS=cuda,cpu python train.py
# or in code:
jax.config.update('jax_platforms', ('cpu', 'cuda'))
Defensive patterns

Strategy: validation

Validate before calling

import jax
try:
    jax.devices('cpu')
except RuntimeError:
    jax.config.update('jax_platforms', None)  # or ('cpu', <accel>)
# now safe to use pure_callback

Try / catch

try:
    jax.jit(f_with_pure_callback)(x)
except RuntimeError as e:
    if 'failed to find a local CPU device' in str(e):
        os.environ['JAX_PLATFORMS'] = 'cpu,cuda'  # restart process
    raise

Prevention

When it happens

Trigger: Running jax.pure_callback (or io_callback with the same impl path) while jax_platforms / JAX_PLATFORMS excludes 'cpu' (e.g. JAX_PLATFORMS=cuda), or an environment where the CPU client failed to initialize.

Common situations: Setting JAX_PLATFORMS=tpu or =cuda for deterministic device selection and then using pure_callback; multi-process GPU jobs that intentionally hide the CPU backend; upgrading JAX where platform restriction became stricter.

Understand the failure class

Background: "environment variable is not set" and "Missing keys in environment" errors: what missing required env var messages mean and how to fix them — this error's family across 28 libraries.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/dd0da9283049e666. Report an issue: GitHub.