{"record":{"id":"b8a2932a9458b50b","repo":"jax-ml/jax","slug":"lookahead-is-not-supported-for-xla-pipeline-emitte","errorCode":null,"errorMessage":"Lookahead is not supported for XLA pipeline emitter lowering.","messagePattern":"Lookahead is not supported for XLA pipeline emitter lowering\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":1271,"sourceCode":"          raise NotImplementedError(\n              \"All block dimensions must be Elements or none of them can be\"\n              \" Elements.\"\n          )\n        padding = [\n            bd.padding if isinstance(bd, pallas_core.Element) else (0, 0)\n            for bd in bm.block_shape\n        ]\n        pad_low, pad_high = map(list, zip(*padding))\n        block_params[\"window_kind\"] = ir.Attribute.parse(\n            f\"#tpu.element_window<{pad_low},{pad_high}>\"\n        )\n      if pipeline_mode is not None:\n        if not isinstance(pipeline_mode, pallas_core.Buffered):\n          raise LoweringException(\n              f\"Unsupported pipeline mode: {pipeline_mode}.\"\n          )\n        if pipeline_mode.use_lookahead:\n          raise NotImplementedError(\n              \"Lookahead is not supported for XLA pipeline emitter lowering.\"\n          )\n        buffer_count = pipeline_mode.buffer_count\n        if buffer_count < 1 or buffer_count > 2:\n          raise LoweringException(\n              \"Only single (1) and double (2) buffering are supported. Got\"\n              f\" {buffer_count}.\"\n          )\n        pipeline_mode_str = \"synchronous\" if buffer_count == 1 else \"double_buffered\"\n        block_params[\"pipeline_mode\"] = ir.Attribute.parse(\n            f\"#tpu.pipeline_mode<{pipeline_mode_str}>\"\n        )\n        if pipeline_mode.revisit is not None:\n          if (\n              pipeline_mode.revisit == pallas_core.RevisitMode.ANY\n              and buffer_count > 1\n          ):\n            raise LoweringException(","sourceCodeStart":1253,"sourceCodeEnd":1289,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L1253-L1289","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Set use_lookahead=False (or omit it) in the Buffered pipeline mode","Rely on double buffering (buffer_count=2) for overlap instead of lookahead","If lookahead pipelining is required, check for an alternative lowering path or a JAX version that supports it on TPU"],"exampleFix":"# before\npipeline_mode=pallas_core.Buffered(buffer_count=2, use_lookahead=True)\n\n# after\npipeline_mode=pallas_core.Buffered(buffer_count=2, use_lookahead=False)","handlingStrategy":"validation","validationCode":"from jax._src.pallas import pallas_core\ndef check_no_lookahead(pm):\n    if isinstance(pm, pallas_core.Buffered) and pm.use_lookahead:\n        raise ValueError('lookahead pipelining unsupported by XLA emitter lowering')","typeGuard":"def is_emitter_compatible(pm) -> bool:\n    from jax._src.pallas import pallas_core\n    return not (isinstance(pm, pallas_core.Buffered) and pm.use_lookahead)","tryCatchPattern":"try:\n    lower_jaxpr_into_pipelined_module(..., pipeline_mode=pm)\nexcept NotImplementedError as e:\n    if 'Lookahead' in str(e):\n        pm = dataclasses.replace(pm, use_lookahead=False); retry","preventionTips":["Leave use_lookahead at its default False","Use double buffering for DMA overlap instead of lookahead"],"tags":["jax","pallas","tpu","pipelining","lookahead"],"backgroundTag":"pallas-unsupported-lookahead","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}