{"record":{"id":"dd0da9283049e666","repo":"jax-ml/jax","slug":"jax-pure-callback-failed-to-find-a-local-cpu-devic","errorCode":null,"errorMessage":"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.","messagePattern":"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\\.","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"jax/_src/callback.py","lineNumber":85,"sourceCode":"  in_tree: tree_util.PyTreeDef  # (args, kwargs) pytree for `callback_func`.\n\n  def __call__(self, *flat_args: Array) -> Sequence[Array]:\n    args, kwargs = tree_util.tree_unflatten(self.in_tree, flat_args)\n    return tree_util.tree_leaves(self.callback_func(*args, **kwargs))\n\n\ndef pure_callback_impl(\n    *args,\n    result_avals,\n    callback: _FlatCallback,\n    sharding: Sharding | None,\n    vmap_method: str | None,\n):\n  del sharding, vmap_method, result_avals\n  try:\n    cpu_device, *_ = xb.local_devices(backend=\"cpu\")\n  except RuntimeError as e:\n    raise RuntimeError(\n        \"jax.pure_callback failed to find a local CPU device to place the\"\n        \" inputs on. Make sure \\\"cpu\\\" is listed in --jax_platforms or the\"\n        \" JAX_PLATFORMS environment variable.\"\n    ) from e\n  args = api.device_put(args, cpu_device)\n  with config.default_device(cpu_device):\n    try:\n      return tree_util.tree_map(np.asarray, callback(*args))\n    except BaseException:\n      logger.exception(\"jax.pure_callback failed\")\n      raise\n\n\npure_callback_p.def_impl(functools.partial(dispatch.apply_primitive,\n                                           pure_callback_p))\n\n\n@pure_callback_p.def_abstract_eval","sourceCodeStart":67,"sourceCodeEnd":103,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/callback.py#L67-L103","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Add 'cpu' to the platform list: JAX_PLATFORMS=cpu,cuda or jax.config.update('jax_platforms', ('cpu','cuda'))","Unset JAX_PLATFORMS entirely if CPU restriction is not required","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"],"exampleFix":"# before\nJAX_PLATFORMS=cuda python train.py  # pure_callback -> RuntimeError\n\n# after\nJAX_PLATFORMS=cuda,cpu python train.py\n# or in code:\njax.config.update('jax_platforms', ('cpu', 'cuda'))","handlingStrategy":"validation","validationCode":"import jax\ntry:\n    jax.devices('cpu')\nexcept RuntimeError:\n    jax.config.update('jax_platforms', None)  # or ('cpu', <accel>)\n# now safe to use pure_callback","typeGuard":null,"tryCatchPattern":"try:\n    jax.jit(f_with_pure_callback)(x)\nexcept RuntimeError as e:\n    if 'failed to find a local CPU device' in str(e):\n        os.environ['JAX_PLATFORMS'] = 'cpu,cuda'  # restart process\n    raise","preventionTips":["Never set JAX_PLATFORMS without including 'cpu' if pure_callback is used","Prefer jax.config.update('jax_platforms', ('cpu', accelerator)) in entry scripts","Add a startup assertion that jax.devices('cpu') succeeds"],"tags":["jax","environment","pure-callback","cpu-device","jax-platforms"],"backgroundTag":"missing-env-var","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}