jax-ml/jax · error · NotImplementedError

Lookahead is not supported for XLA pipeline emitter lowering

Error message

Lookahead is not supported for XLA pipeline emitter lowering.

What it means

pallas_core.Buffered has a use_lookahead option (software pipelining with lookahead), but the XLA pipeline emitter lowering path does not implement it. If use_lookahead=True reaches the TPU lowering, it raises NotImplementedError.

Source

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

          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 (
              pipeline_mode.revisit == pallas_core.RevisitMode.ANY
              and buffer_count > 1
          ):
            raise LoweringException(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Set use_lookahead=False (or omit it) in the Buffered pipeline mode
  2. Rely on double buffering (buffer_count=2) for overlap instead of lookahead
  3. If lookahead pipelining is required, check for an alternative lowering path or a JAX version that supports it on TPU

Example fix

# before
pipeline_mode=pallas_core.Buffered(buffer_count=2, use_lookahead=True)

# after
pipeline_mode=pallas_core.Buffered(buffer_count=2, use_lookahead=False)
Defensive patterns

Strategy: validation

Validate before calling

from jax._src.pallas import pallas_core
def check_no_lookahead(pm):
    if isinstance(pm, pallas_core.Buffered) and pm.use_lookahead:
        raise ValueError('lookahead pipelining unsupported by XLA emitter lowering')

Type guard

def is_emitter_compatible(pm) -> bool:
    from jax._src.pallas import pallas_core
    return not (isinstance(pm, pallas_core.Buffered) and pm.use_lookahead)

Try / catch

try:
    lower_jaxpr_into_pipelined_module(..., pipeline_mode=pm)
except NotImplementedError as e:
    if 'Lookahead' in str(e):
        pm = dataclasses.replace(pm, use_lookahead=False); retry

Prevention

When it happens

Trigger: Passing pipeline_mode=pallas_core.Buffered(buffer_count=..., use_lookahead=True) (or using an API/experimental flag that sets lookahead) when lowering through the XLA pipeline emitter on TPU.

Common situations: Enabling lookahead software pipelining expecting better DMA overlap; using a newer Pallas feature flag with an older lowering path; kernels tuned with lookahead on a different backend then run through the XLA emitter.

Related errors

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