jax-ml/jax · error · LoweringException

Only single (1) and double (2) buffering are supported. Got

Error message

Only single (1) and double (2) buffering are supported. Got {buffer_count}.

What it means

The TPU lowering supports only single buffering (buffer_count=1, synchronous) and double buffering (buffer_count=2) for pipelined DMA. Any other buffer_count (0, 3, ...) raises LoweringException, because the emitted #tpu.pipeline_mode attribute can only be 'synchronous' or 'double_buffered'.

Source

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

            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 (
              pipeline_mode.revisit == pallas_core.RevisitMode.ANY
              and buffer_count > 1
          ):
            raise LoweringException(
                "RevisitMode.ANY is not supported with double buffering."
            )
          block_params["revisit_mode"] = ir.Attribute.parse(
              f"#tpu.revisit_mode<{pipeline_mode.revisit.value}>"
          )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use buffer_count=2 (double buffering), the maximum supported
Use buffer_count=1 for synchronous, memory-frugal pipelines
Restructure the kernel (e.g. smaller tiles) instead of adding more buffers to hide latency

Example fix

# before
pipeline_mode=pallas_core.Buffered(buffer_count=3)

# after
pipeline_mode=pallas_core.Buffered(buffer_count=2)

Defensive patterns

Strategy: validation

Validate before calling

from jax._src.pallas import pallas_core
def check_buffer_count(pm):
    if isinstance(pm, pallas_core.Buffered) and not (1 <= pm.buffer_count <= 2):
        raise ValueError('buffer_count must be 1 or 2 on TPU')

Type guard

def has_supported_buffer_count(pm) -> bool:
    from jax._src.pallas import pallas_core
    return not isinstance(pm, pallas_core.Buffered) or pm.buffer_count in (1, 2)

Try / catch

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

Prevention

Clamp buffer_count to 1 or 2 before passing pipeline_mode
Use smaller tiles rather than more buffers to tune performance

When it happens

Trigger: Passing pipeline_mode=pallas_core.Buffered(buffer_count=3) (or 0) to the TPU lowering; requesting triple buffering for deeper DMA overlap.

Common situations: Tuning kernels by increasing buffer counts expecting more prefetch depth; porting GPU-style multi-stage pipelining to Pallas TPU; scripts parameterizing buffer_count without clamping.

Related errors

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