jax-ml/jax · error · NotImplementedError

accessing .backend in multi-lowering setting. This can occur

Error message

accessing .backend in multi-lowering setting. This can occur when lowering a primitive that has not been adapted to multi-platform lowering

What it means

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.

Source

Thrown at jax/_src/interpreters/mlir.py:894

    self.channel_iterator = channel_iterator
    self.keepalives = keepalives
    self.host_callbacks = host_callbacks
    self.shape_poly_state = (
      shape_poly_state or ShapePolyLoweringState((), tuple(platforms)))
    self.all_default_mem_kind = all_default_mem_kind
    self.lowering_parameters = lowering_parameters
    self.sharding_attr_cache = ({} if sharding_attr_cache is None else sharding_attr_cache)
    self.aval_to_ir_types_cache = ({} if aval_to_ir_types_cache is None else aval_to_ir_types_cache)
    self.pallas_lowering_cache = ({} if pallas_lowering_cache is None else pallas_lowering_cache)
    self.pallas_collective_id_mapping = (CollectiveIdMapping()
                                         if pallas_collective_id_mapping is None
                                         else pallas_collective_id_mapping)

  def get_backend(self, optional: bool = False) -> xc.Client | None:
    if len(self.platforms) > 1:
      if optional:
        return None
      raise NotImplementedError(
        "accessing .backend in multi-lowering setting. This can occur when "
        "lowering a primitive that has not been adapted to multi-platform "
        "lowering")
    if self.backend is not None:
      if xb.canonicalize_platform(self.backend.platform) != self.platforms[0]:
        if optional:
          return None
        raise ValueError(
          "the platform for the specified backend "
          f"{xb.canonicalize_platform(self.backend.platform)} is different "
          f"from the lowering platform {self.platforms[0]}")
      return self.backend
    return xb.get_backend(self.platforms[0])

  def new_channel(self) -> int:
    channel = next(self.channel_iterator)
    # `xla::HostCallback` requires a 16-bit channel ID.
    if channel >= (1 << 16):

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

# before
backend = ctx.get_backend()  # raises in multi-lowering

# after
backend = ctx.get_backend(optional=True)
if backend is None:
    # platform-agnostic lowering path
    ...

Defensive patterns

Strategy: fallback

Validate before calling

if len(ctx.platforms) > 1:
    backend = ctx.get_backend(optional=True)  # returns None, no raise
else:
    backend = ctx.get_backend()

Try / catch

try:
    backend = ctx.get_backend()
except NotImplementedError:
    backend = ctx.get_backend(optional=True)
    if backend is None:
        # platform-agnostic path
        ...

Prevention

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

When it happens

Trigger: 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.

Common situations: 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.

Related errors

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