{"record":{"id":"c2cb8cbe8eb13c22","repo":"jax-ml/jax","slug":"accessing-backend-in-multi-lowering-setting-this","errorCode":null,"errorMessage":"accessing .backend in multi-lowering setting. This can occur when lowering a primitive that has not been adapted to multi-platform lowering","messagePattern":"accessing \\.backend in multi-lowering setting\\. This can occur when lowering a primitive that has not been adapted to multi-platform lowering","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/mlir.py","lineNumber":894,"sourceCode":"    self.channel_iterator = channel_iterator\n    self.keepalives = keepalives\n    self.host_callbacks = host_callbacks\n    self.shape_poly_state = (\n      shape_poly_state or ShapePolyLoweringState((), tuple(platforms)))\n    self.all_default_mem_kind = all_default_mem_kind\n    self.lowering_parameters = lowering_parameters\n    self.sharding_attr_cache = ({} if sharding_attr_cache is None else sharding_attr_cache)\n    self.aval_to_ir_types_cache = ({} if aval_to_ir_types_cache is None else aval_to_ir_types_cache)\n    self.pallas_lowering_cache = ({} if pallas_lowering_cache is None else pallas_lowering_cache)\n    self.pallas_collective_id_mapping = (CollectiveIdMapping()\n                                         if pallas_collective_id_mapping is None\n                                         else pallas_collective_id_mapping)\n\n  def get_backend(self, optional: bool = False) -> xc.Client | None:\n    if len(self.platforms) > 1:\n      if optional:\n        return None\n      raise NotImplementedError(\n        \"accessing .backend in multi-lowering setting. This can occur when \"\n        \"lowering a primitive that has not been adapted to multi-platform \"\n        \"lowering\")\n    if self.backend is not None:\n      if xb.canonicalize_platform(self.backend.platform) != self.platforms[0]:\n        if optional:\n          return None\n        raise ValueError(\n          \"the platform for the specified backend \"\n          f\"{xb.canonicalize_platform(self.backend.platform)} is different \"\n          f\"from the lowering platform {self.platforms[0]}\")\n      return self.backend\n    return xb.get_backend(self.platforms[0])\n\n  def new_channel(self) -> int:\n    channel = next(self.channel_iterator)\n    # `xla::HostCallback` requires a 16-bit channel ID.\n    if channel >= (1 << 16):","sourceCodeStart":876,"sourceCodeEnd":912,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/mlir.py#L876-L912","documentation":"LoweringRuleContext.get_backend() is ambiguous when the lowering targets multiple platforms at once. JAX's multi-platform lowering pipeline does not pin a single backend, so primitives that reach for ctx.backend in that setting are unsupported and raise NotImplementedError, with optional=True as the graceful (returns None) alternative.","triggerScenarios":"A primitive's lowering rule calls ctx.get_backend() while the computation is being lowered for multiple platforms simultaneously (e.g. via multi-platform export/compile paths); typically affects JAX-internal or custom primitives not yet adapted to the newer multi-platform API.","commonSituations":"Upgrading JAX to a version where multi-platform lowering became the default for certain paths; custom primitives written against the old ctx.backend API; using export/serialization features that lower for CPU+GPU together.","solutions":["Upgrade JAX — the primitive may have been adapted to multi-platform lowering in a newer release","In your lowering rule, call ctx.get_backend(optional=True) and handle None, or use platform-specific logic via ctx.platforms","Re-run the computation pinned to a single platform so the lowering is not multi-platform"],"exampleFix":"# before\nbackend = ctx.get_backend()  # raises in multi-lowering\n\n# after\nbackend = ctx.get_backend(optional=True)\nif backend is None:\n    # platform-agnostic lowering path\n    ...","handlingStrategy":"fallback","validationCode":"if len(ctx.platforms) > 1:\n    backend = ctx.get_backend(optional=True)  # returns None, no raise\nelse:\n    backend = ctx.get_backend()","typeGuard":null,"tryCatchPattern":"try:\n    backend = ctx.get_backend()\nexcept NotImplementedError:\n    backend = ctx.get_backend(optional=True)\n    if backend is None:\n        # platform-agnostic path\n        ...","preventionTips":["In custom primitives always call get_backend(optional=True) first","Pin compilation to one platform when a primitive needs ctx.backend","Track JAX release notes for multi-platform lowering migration guides"],"tags":["jax","lowering","multi-platform","primitive"],"backgroundTag":"api-version-incompatibility","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}