{"record":{"id":"09641a6a0883a8fa","repo":"jax-ml/jax","slug":"grid-must-be-specified-when-using-lookahead","errorCode":null,"errorMessage":"Grid must be specified when using lookahead.","messagePattern":"Grid must be specified when using lookahead\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/pipeline.py","lineNumber":1489,"sourceCode":"    in_specs and out_specs.\n  \"\"\"\n  # TODO(levskaya): generalize argument tree handling here and in emit_pipeline.\n  num_in_specs = len(in_specs)\n  in_specs = _normalize_specs(in_specs)\n  out_specs = _normalize_specs(out_specs)\n  in_specs, in_refs = _filter_specs_and_refs(in_specs, refs[:num_in_specs])\n  out_specs, out_refs = _filter_specs_and_refs(out_specs, refs[num_in_specs:])\n  def make_input_bref(in_spec, in_ref):\n    in_aval = _ref_to_value_aval(in_ref)\n    buffer_count = 2\n    use_lookahead = False\n    prefetched_count = 0\n    if has_buffering := in_spec.pipeline_mode is not None:\n      buffer_count = in_spec.pipeline_mode.buffer_count\n      use_lookahead = in_spec.pipeline_mode.use_lookahead\n      prefetched_count = in_spec.pipeline_mode.prefetched_count\n    if use_lookahead and grid is None:\n      raise ValueError(\"Grid must be specified when using lookahead.\")\n    is_trivial = _spec_has_trivial_windowing(in_spec, grid, in_aval.shape)\n    if not has_buffering and is_trivial:\n      buffer_count = 1\n\n    sms = (in_ref.memory_space if isinstance(in_ref, state.TransformedRef) else\n           core.typeof(in_ref).memory_space)\n    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)","sourceCodeStart":1471,"sourceCodeEnd":1507,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/pipeline.py#L1471-L1507","documentation":"When an input spec uses pipeline_mode with use_lookahead=True, the pipeline must know the kernel grid to compute fetch indices ahead of time. If grid is None (dynamic/deferred grid) at BufferSpec binding time, this ValueError is raised.","triggerScenarios":"Setting PipelineMode(use_lookahead=True) in an in_spec while calling the pipeline entry point without a concrete grid (grid=None), e.g. relying on a later grid specification.","commonSituations":"Adapting a kernel that previously ran without lookahead and without an explicit grid; using dynamic grid specification features while enabling lookahead.","solutions":["Pass an explicit tuple of Python ints / JAX arrays as grid when using lookahead","Or disable lookahead in the pipeline_mode of that input spec","Verify grid is passed at the same call site that receives in_specs with pipeline_mode"],"exampleFix":"# before\nemit_pipeline(..., grid=None)  # with use_lookahead=True\n# after\nemit_pipeline(..., grid=(num_steps,))","handlingStrategy":"validation","validationCode":"assert grid is not None or not any(\n    getattr(s, 'pipeline_mode', None) and s.pipeline_mode.use_lookahead for s in in_specs), \\\n    'lookahead requires an explicit grid'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always pass a concrete grid tuple when any spec uses use_lookahead","Compute grid before building specs so it can't be None"],"tags":["jax","pallas","pipeline","lookahead","grid"],"backgroundTag":"missing-required-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}