{"record":{"id":"6adfb45849cfb7e0","repo":"jax-ml/jax","slug":"partitioned-callback-not-implemented-on-platform","errorCode":null,"errorMessage":"Partitioned callback not implemented on {platform} backend.","messagePattern":"Partitioned callback not implemented on (.+?) backend\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/callback.py","lineNumber":805,"sourceCode":"    has_side_effect: Whether the callback has side effects.\n    returns_token: Whether the callback should return a token.\n    partitioned: If True, then `callback` is called on local shards only. If\n      False, then `callback` is called on all shards.\n    sharding: The sharding of the callback.\n\n  Returns:\n    A tuple of MLIR result values, a new token (if any), and the host callback\n    object.\n  \"\"\"\n  if len(ctx.module_context.platforms) > 1:\n    raise NotImplementedError(\"multi-platform lowering for python_callback\")\n  platform = ctx.module_context.platforms[0]\n  if platform not in {\"cpu\", \"cuda\", \"rocm\", \"tpu\", \"oneapi\"}:\n    raise ValueError(\n        f\"`EmitPythonCallback` not supported on {platform} backend.\")\n  if partitioned:\n    if platform not in {\"cpu\", \"cuda\", \"rocm\", \"oneapi\"}:\n      raise NotImplementedError(\n          f\"Partitioned callback not implemented on {platform} backend.\")\n    if result_avals:\n      raise ValueError(\"Partitioned callback not supported with return values.\")\n  backend: xc.Client = cast(xc.Client, ctx.module_context.get_backend())\n  result_shapes = [_aval_to_xla_shape(aval) for aval in result_avals]\n  operand_shapes = [_aval_to_xla_shape(aval) for aval in operand_avals]\n\n  # First we apply checks to ensure output shapes and dtypes match the expected\n  # ones.\n  def _wrapped_callback(*args):\n    out_vals = callback(*args)\n    if len(out_vals) != len(result_avals):\n      raise RuntimeError(\n          \"Mismatched number of outputs from callback. \"\n          \"Expected: {}, Actual: {}\".format(len(result_avals), len(out_vals)))\n    # Handle Python literals, and custom arrays, e.g., tf.Tensor.\n    out_vals = tuple(dtypes.canonicalize_value(np.asarray(a)) for a in out_vals)\n    for i, (out_val, out_aval) in enumerate(zip(out_vals, result_avals)):","sourceCodeStart":787,"sourceCodeEnd":823,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/callback.py#L787-L823","documentation":"Raised during lowering of a partitioned JAX callback (io_callback/pure_callback with _partitioned=True) on a backend that doesn't support partitioning. Only cpu, cuda, rocm, and oneapi support partitioned callbacks; TPU does not.","triggerScenarios":"Calling io_callback(..., _partitioned=True) (as done inside jax.debug, pallas, or TPU-scheduled code) while compiling/executing on the TPU backend (or any platform outside the supported set).","commonSituations":"Running code that was written for GPU partitioned callbacks on a TPU device or cloud TPU VM; using libraries (e.g. extendedblas, pallas TPU) that internally request partitioned callbacks.","solutions":["Switch to a non-partitioned callback: drop _partitioned=True and use a plain jax.experimental.io_callback / pure_callback","Run on a supported backend (cpu/cuda/rocm/oneapi) if partitioning is required","Upgrade JAX — check release notes for added TPU partitioned-callback support"],"exampleFix":"# before\nio_callback(fn, result_avals, *args, _partitioned=True)  # on TPU\n# after\nio_callback(fn, result_avals, *args)  # unpartitioned, works on TPU","handlingStrategy":"validation","validationCode":"from jax._src.callback import _PARTITIONED_PLATFORMS  # illustrative\nsupported = {'cpu','cuda','rocm','oneapi'}\nplatform = jax.default_backend()\nif platform not in supported:\n    callback_kwargs.pop('_partitioned', None)  # degrade gracefully","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Gate _partitioned=True behind jax.default_backend() checks","Document backend requirements when using partitioned callbacks"],"tags":["jax","callback","tpu","sharding","backend-support"],"backgroundTag":"unsupported-platform-feature","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}