jax-ml/jax · error · LoweringException

Unsupported pipeline mode: {pipeline_mode}.

Error message

Unsupported pipeline mode: {pipeline_mode}.

What it means

The pipeline_mode argument of the TPU lowering must be a pallas_core.Buffered instance (single or double buffering). Any other value — an int, string, or a newer/unrecognized pipeline mode object — raises LoweringException with the offending value.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:1267

            isinstance(bd, (pallas_core.Element, pallas_core.Squeezed))
            for bd in bm.block_shape
        ]
        if not all(is_element_or_squeezed_block):
          raise NotImplementedError(
              "All block dimensions must be Elements or none of them can be"
              " Elements."
          )
        padding = [
            bd.padding if isinstance(bd, pallas_core.Element) else (0, 0)
            for bd in bm.block_shape
        ]
        pad_low, pad_high = map(list, zip(*padding))
        block_params["window_kind"] = ir.Attribute.parse(
            f"#tpu.element_window<{pad_low},{pad_high}>"
        )
      if pipeline_mode is not None:
        if not isinstance(pipeline_mode, pallas_core.Buffered):
          raise LoweringException(
              f"Unsupported pipeline mode: {pipeline_mode}."
          )
        if pipeline_mode.use_lookahead:
          raise NotImplementedError(
              "Lookahead is not supported for XLA pipeline emitter lowering."
          )
        buffer_count = pipeline_mode.buffer_count
        if buffer_count < 1 or buffer_count > 2:
          raise LoweringException(
              "Only single (1) and double (2) buffering are supported. Got"
              f" {buffer_count}."
          )
        pipeline_mode_str = "synchronous" if buffer_count == 1 else "double_buffered"
        block_params["pipeline_mode"] = ir.Attribute.parse(
            f"#tpu.pipeline_mode<{pipeline_mode_str}>"
        )
        if pipeline_mode.revisit is not None:
          if (

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass pallas_core.Buffered(buffer_count=1 or 2) (or None) for pipeline_mode
2. Prefer the public pallas_call API, which constructs pipeline_mode correctly, over calling the lowering directly
3. Check the signature of lower_jaxpr_into_pipelined_module in your installed JAX version and update call sites

Example fix

# before
lower_jaxpr_into_pipelined_module(..., pipeline_mode=2)

# after
from jax._src.pallas import pallas_core
lower_jaxpr_into_pipelined_module(..., pipeline_mode=pallas_core.Buffered(buffer_count=2))

from jax._src.pallas import pallas_core
def valid_pipeline_mode(pm):
    return pm is None or isinstance(pm, pallas_core.Buffered)

def is_supported_pipeline_mode(pm) -> bool:
    from jax._src.pallas import pallas_core
    return pm is None or isinstance(pm, pallas_core.Buffered)

try:
    lower_jaxpr_into_pipelined_module(..., pipeline_mode=pm)
except LoweringException as e:
    if 'Unsupported pipeline mode' in str(e):
        pm = pallas_core.Buffered(buffer_count=2); retry

When it happens

Trigger: Passing pipeline_mode as a non-Buffered value to lower_jaxpr_into_pipelined_module / lower_jaxpr_to_pipelined_module, e.g. pipeline_mode=2 or pipeline_mode='double_buffered' instead of pallas_core.Buffered(buffer_count=2).

Common situations: Calling the lowering API directly instead of via pallas_call; code written against an older/newer JAX where the pipeline_mode representation changed (e.g. from ints/enums to Buffered dataclass); copying internal example code with a stale signature.

Related errors

• Unsupported block dimension type: {type(bd)} for block shape
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Acc ref must be at least 2D, got shape {shape}
• Acc ref dtype must be float32 or int32, got {dtype}

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