jax-ml/jax · error · ValueError

Only interpret mode is supported on CPU backend.

Error message

Only interpret mode is supported on CPU backend.

What it means

The Pallas CPU backend only implements interpret mode — there is no real Pallas compilation pipeline for CPU. If a pallas_call is lowered on CPU without interpret=True, the registered cpu_lowering rule raises ValueError('Only interpret mode is supported on CPU Backend.') at compile time.

Source

Thrown at jax/_src/pallas/pallas_call.py:885

      from jax._src.pallas.mosaic_gpu.interpret import params as gpu_params  # pyrefly: ignore[missing-import]
    except ImportError:
      pass
    else:
      if isinstance(interpret, gpu_params.InterpretGPUParams):
        impl = partial(
            mosaic_gpu_interpret.interpret_pallas_call,
            interpret_params=interpret,
            **params,
        )

    return mlir.lower_fun(impl, multiple_results=True)(ctx, *in_nodes)

  def cpu_lowering(
      ctx: mlir.LoweringRuleContext,
      *in_nodes: ir.Value | Sequence[ir.Value],
      **params,
  ):
    raise ValueError("Only interpret mode is supported on CPU backend.")

  def tpu_lowering(
      ctx: mlir.LoweringRuleContext,
      *in_nodes: ir.Value | Sequence[ir.Value],
      **params,
  ):
    compiler_params = params.get("compiler_params")
    if compiler_params is not None:
      rule = pallas_core.get_lowering_rule(type(compiler_params), "tpu")
      if rule is not None:
        return rule(ctx, *in_nodes, **params)

    try:
      from jax._src.pallas.mosaic import pallas_call_registration as mosaic_tpu_backend  # pyrefly: ignore[missing-import]
    except ImportError:
      raise _unsupported_lowering_error("tpu") from None

    return mosaic_tpu_backend.pallas_call_tpu_lowering_rule(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Set interpret=True in pallas_call for CPU runs (debugging mode; slow but functional)
Run the code on an actual GPU/TPU backend where the kernel can be compiled
Gate hardware-specific code: use interpret=True when jax.devices()[0].platform == 'cpu'
Use interpreter mode in CI and real hardware in production

Example fix

// before
out = pallas_call(kernel, grid, out_shape)(x, y)  # fails on CPU
// after
import jax
interpret = jax.devices()[0].platform == 'cpu'
out = pallas_call(kernel, grid, out_shape, interpret=interpret)(x, y)

Defensive patterns

Strategy: validation

Validate before calling

import jax
is_cpu = jax.devices()[0].platform == 'cpu'
if is_cpu and not interpret:
    interpret = True  # or raise early with a clear message

Type guard

def needs_interpret(interpret: bool) -> bool:
    return jax.devices()[0].platform == 'cpu' and not interpret

Try / catch

try:
    pallas_call(kernel, grid, out_shape)(x)
except ValueError as e:
    if 'interpret mode' in str(e):
        out = pallas_call(kernel, grid, out_shape, interpret=True)(x)

Prevention

Parameterize interpret by device platform in all Pallas code
Use interpret=True for local CPU development and CI
Verify jax.devices() platform before compiling kernels

When it happens

Trigger: Running pallas_call(..., interpret=False) (the default) on a CPU-only machine or with JAX_PLATFORMS=cpu; writing kernels for TPU/GPU and executing/testing locally on CPU.

Common situations: Developers without GPU/TPU hardware running Pallas examples locally; CI machines that are CPU-only; accidentally forcing the CPU backend via jax.config or JAX_PLATFORM_NAME while developing Triton/Mosaic kernels.

Related errors

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