jax-ml/jax · error · ValueError

Compiler params for platform {platform} cannot be used for {

Error message

Compiler params for platform {platform} cannot be used for {expected_platform} lowering.

What it means

Pallas registers one lowering rule per compiler-params class per platform. get_lowering_rule looks up the rule for e.g. TPUCache or TritonCompilerParams and raises if the registered platform (TPU vs GPU) does not match the platform the computation is being lowered for. It guards against applying TPU kernel params to a GPU lowering and vice versa.

Source

Thrown at jax/_src/pallas/core.py:186

  do nothing.
  """
  return checkify.debug_check(condition, message)


_backend_lowering_rules = {}


def register_lowering_rule(params_cls, rule, platform: str):
  _backend_lowering_rules[params_cls] = (rule, platform)


def get_lowering_rule(params_cls, expected_platform: str):
  rule_info = _backend_lowering_rules.get(params_cls)
  if rule_info is None:
    return None
  rule, platform = rule_info
  if platform != expected_platform:
    raise ValueError(
        f"Compiler params for platform {platform} cannot be used for"
        f" {expected_platform} lowering."
    )
  return rule


@enum.unique
class RevisitMode(enum.Enum):
  """Specifies whether an output buffer supports revisiting.

  By default, buffers can only be safely revisited at the next iteration
  (immediate revisiting). If revisited at any other iteration, the buffer state
  should be considered undefined.

  If revisiting at any arbitrary iteration is required, use RevisitMode.ANY.
  This will insert additional DMAs as needed to restore the buffer state.

  Input buffers ignore revisit mode: as inputs read data from memory, their

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Match the lowering decorator to the backend you run on: use triton/gpu_lowering on GPU, tpu_lowering on TPU
Make the kernel import conditional on jax.default_backend()
Skip or substitute CPU reference implementations when running on unsupported backends

Example fix

# before
@pl.tpu_lowering(compile_params)
def kernel(...): ...
# run on GPU -> error

# after
from jax.experimental.pallas import tpu, triton
if jax.default_backend() == 'tpu':
    kernel = pl.pallas_call(kernel, ...).lowering(tpu.lowering(...))
else:
    kernel = triton.pallas_call(kernel, ...)

Defensive patterns

Strategy: type-guard

Validate before calling

import jax
backend = jax.default_backend()
# choose lowering to match backend before building the kernel

Type guard

def pick_lowering(backend):
    if backend == 'tpu':
        import jax.experimental.pallas.tpu as tpu
        return tpu.lowering
    import jax.experimental.pallas.triton as triton
    return triton.pallas_call

Try / catch

catch ValueError from lowering and re-dispatch to the correct backend kernel

Prevention

Branch kernel imports on jax.default_backend()
Run CI on the same backend as production for Pallas kernels

When it happens

Trigger: Decorating a kernel with @pl.tpu_lowering(...) (or gpu_lowering/Triton) and then executing/lowering it on the other backend — e.g. writing a Triton kernel via jax.experimental.pallas.triton but running on TPU, or using TPU compiler params on GPU/CPU.

Common situations: Developing Pallas kernels on GPU then running on TPU (or CI on CPU); importing a pallas.triton module in a TPU job; version changes that moved decorators between modules.

Related errors

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