{"record":{"id":"918681356f56eece","repo":"jax-ml/jax","slug":"output-buffering-does-not-support-lookahead","errorCode":null,"errorMessage":"Output buffering does not support lookahead.","messagePattern":"Output buffering does not support lookahead\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/pipeline.py","lineNumber":1514,"sourceCode":"    return BufferedRef.input(\n        in_spec,\n        in_aval,\n        buffer_count,\n        grid_rank=len(grid),\n        use_lookahead=use_lookahead,\n        source_memory_space=sms,\n        tiling=tiling,\n        is_trivial_windowing=is_trivial,\n        prefetched_count=prefetched_count,\n    )\n  in_brefs = jax.tree.map(make_input_bref, in_specs, in_refs)\n  def make_output_bref(out_spec, out_ref):\n    out_aval = _ref_to_value_aval(out_ref)\n    buffer_count = 2\n    if has_buffering := out_spec.pipeline_mode is not None:\n      buffer_count = out_spec.pipeline_mode.buffer_count\n      if out_spec.pipeline_mode.use_lookahead:\n        raise ValueError(\"Output buffering does not support lookahead.\")\n    is_trivial = _spec_has_trivial_windowing(out_spec, grid, out_aval.shape)\n    if not has_buffering and is_trivial:\n      buffer_count = 1\n\n    sms = (out_ref.memory_space if isinstance(out_ref, state.TransformedRef)\n           else core.typeof(out_ref).memory_space)\n    return BufferedRef.output(\n        out_spec,\n        out_aval,\n        buffer_count,\n        source_memory_space=sms,\n        tiling=tiling,\n        is_trivial_windowing=is_trivial,\n    )\n  out_brefs = jax.tree.map(make_output_bref, out_specs, out_refs)\n  return (*in_brefs, *out_brefs)\n\ndef _resolve_core_info(core_axis: tuple[int | str, ...] | int | str | None):","sourceCodeStart":1496,"sourceCodeEnd":1532,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/pipeline.py#L1496-L1532","documentation":"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.","triggerScenarios":"Setting PipelineMode(use_lookahead=True) on an out_spec (output BlockSpec) of a pipelined kernel.","commonSituations":"Applying the same PipelineMode object intended for inputs to all specs via a shared list; copy-pasting in_specs configuration into out_specs.","solutions":["Remove use_lookahead from output specs; keep lookahead only on input specs","If outputs also need double buffering, use plain pipeline_mode buffer_count > 1 without lookahead","Construct in_specs and out_specs separately rather than sharing one spec list"],"exampleFix":"# before\npm = PipelineMode(use_lookahead=True)\nin_specs = [BlockSpec(..., pipeline_mode=pm)]\nout_specs = [BlockSpec(..., pipeline_mode=pm)]\n# after\nin_specs = [BlockSpec(..., pipeline_mode=PipelineMode(use_lookahead=True))]\nout_specs = [BlockSpec(..., pipeline_mode=PipelineMode())]","handlingStrategy":"validation","validationCode":"for s in out_specs:\n    pm = getattr(s, 'pipeline_mode', None)\n    assert pm is None or not pm.use_lookahead, 'lookahead not supported on outputs'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Build in_specs and out_specs separately; never share a PipelineMode with use_lookahead across both"],"tags":["jax","pallas","pipeline","lookahead","output-buffer"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}