jax-ml/jax · error · ValueError

Output buffering does not support lookahead.

Error message

Output buffering does not support lookahead.

What it means

Lookahead (prefetching future iterations) is implemented only for input buffers, where the pipeline can fetch ahead of the compute loop. Output buffers are written as the body runs, so make_output_bref rejects pipeline_mode with use_lookahead=True with this ValueError.

Source

Thrown at jax/_src/pallas/mosaic/pipeline.py:1514

    return BufferedRef.input(
        in_spec,
        in_aval,
        buffer_count,
        grid_rank=len(grid),
        use_lookahead=use_lookahead,
        source_memory_space=sms,
        tiling=tiling,
        is_trivial_windowing=is_trivial,
        prefetched_count=prefetched_count,
    )
  in_brefs = jax.tree.map(make_input_bref, in_specs, in_refs)
  def make_output_bref(out_spec, out_ref):
    out_aval = _ref_to_value_aval(out_ref)
    buffer_count = 2
    if has_buffering := out_spec.pipeline_mode is not None:
      buffer_count = out_spec.pipeline_mode.buffer_count
      if out_spec.pipeline_mode.use_lookahead:
        raise ValueError("Output buffering does not support lookahead.")
    is_trivial = _spec_has_trivial_windowing(out_spec, grid, out_aval.shape)
    if not has_buffering and is_trivial:
      buffer_count = 1

    sms = (out_ref.memory_space if isinstance(out_ref, state.TransformedRef)
           else core.typeof(out_ref).memory_space)
    return BufferedRef.output(
        out_spec,
        out_aval,
        buffer_count,
        source_memory_space=sms,
        tiling=tiling,
        is_trivial_windowing=is_trivial,
    )
  out_brefs = jax.tree.map(make_output_bref, out_specs, out_refs)
  return (*in_brefs, *out_brefs)

def _resolve_core_info(core_axis: tuple[int | str, ...] | int | str | None):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove use_lookahead from output specs; keep lookahead only on input specs
2. If outputs also need double buffering, use plain pipeline_mode buffer_count > 1 without lookahead
3. Construct in_specs and out_specs separately rather than sharing one spec list

Example fix

# before
pm = PipelineMode(use_lookahead=True)
in_specs = [BlockSpec(..., pipeline_mode=pm)]
out_specs = [BlockSpec(..., pipeline_mode=pm)]
# after
in_specs = [BlockSpec(..., pipeline_mode=PipelineMode(use_lookahead=True))]
out_specs = [BlockSpec(..., pipeline_mode=PipelineMode())]

for s in out_specs:
    pm = getattr(s, 'pipeline_mode', None)
    assert pm is None or not pm.use_lookahead, 'lookahead not supported on outputs'

When it happens

Trigger: Setting PipelineMode(use_lookahead=True) on an out_spec (output BlockSpec) of a pipelined kernel.

Common situations: Applying the same PipelineMode object intended for inputs to all specs via a shared list; copy-pasting in_specs configuration into out_specs.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Grid must be specified when using lookahead.
• Lookahead is not supported for XLA pipeline emitter lowering
• Element with {padding=} is not supported.
• Only one of src or dst can be a BufferedRef.
• One of src or dst must be a BufferedRef.

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