jax-ml/jax · error · ValueError

the platform for the specified backend {xb.canonicalize_plat

Error message

the platform for the specified backend {xb.canonicalize_platform(self.backend.platform)} is different from the lowering platform {self.platforms[0]}

What it means

A LoweringRuleContext can carry an explicit backend; if that backend's canonicalized platform differs from the platform the computation is being lowered for, mixing the two would produce invalid code, so JAX raises a ValueError spelling out both platforms.

Source

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

    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):
      raise RuntimeError(
          "Host callback lowering created too many channels. PjRt does not"
          " support more than 65535 channels")
    return channel

  # Adds an IFRT host callback object to the context. A reference to these
  # callbacks will be provided to IFRT during compilation so it can do things
  # like serialize them and keep them alive.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Ensure the backend passed to the lowering context matches the platform being lowered: xb.get_backend(target_platform)
  2. Rebuild the context rather than reusing one created for a different platform
  3. Canonicalize platform names (xb.canonicalize_platform) when comparing/user input, e.g. 'cuda' vs 'gpu'

Example fix

# before
ctx = ctx.replace(backend=xb.get_backend('cpu'))
lower_for(ctx, platforms=['gpu'])  # ValueError

# after
ctx = ctx.replace(backend=xb.get_backend('gpu'))
lower_for(ctx, platforms=['gpu'])
Defensive patterns

Strategy: validation

Validate before calling

from jax._src import xla_bridge as xb

assert xb.canonicalize_platform(backend.platform) == ctx.platforms[0], \
    f'backend {backend.platform} vs lowering {ctx.platforms[0]}'

Prevention

When it happens

Trigger: Constructing or reusing a LoweringRuleContext whose backend was built for platform A (e.g. 'cpu') while lowering targets platform B (e.g. 'gpu'/'tpu'); commonly from manually built contexts, custom primitives reusing a cached ctx, or specifying devices/backends inconsistently across jax.jit(device=...) and lowering contexts.

Common situations: Mismatched jax devices in multi-GPU/TPU setups; tests that build a CPU context but lower GPU computations; custom export tooling that pins a backend then switches platforms; version changes in platform canonicalization naming.

Related errors

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