{"record":{"id":"084bb9bf6e52027c","repo":"jax-ml/jax","slug":"index-map-function-debug-info-func-src-info-for-084bb9","errorCode":null,"errorMessage":"Index map function {debug_info.func_src_info} for {origin} must not capture constants: {closed_jaxpr.consts}","messagePattern":"Index map function (.+?) for (.+?) must not capture constants: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/core.py","lineNumber":698,"sourceCode":"          if (\n              not isinstance(idx_aval, jax_core.ShapedArray)\n              and not idx_aval.shape\n          ):\n            raise ValueError(\n                \"index_map returned a value of type\"\n                f\" {type(idx_aval)} at position {i} with block dimension\"\n                f\" {bd} when it should be a scalar\"\n            )\n    for i, ov in enumerate(out_avals):\n      if ov.shape or ov.dtype not in [jnp.int32, jnp.int64]:\n        raise ValueError(\n            f\"Index map function {debug_info.func_src_info} for \"\n            f\"{origin} must return integer scalars. Output[{i}] has type \"\n            f\"{ov}.\"\n        )\n\n    if closed_jaxpr.consts and not allow_captured_consts:\n      raise ValueError(\n          f\"Index map function {debug_info.func_src_info} for \"\n          f\"{origin} must not capture constants: {closed_jaxpr.consts}\"\n      )\n\n    mapping = BlockMapping(\n        block_shape=block_shape,\n        transformed_block_aval=block_aval,  # There are no transforms by default\n        index_map_jaxpr=closed_jaxpr,\n        index_map_out_tree=out_avals.tree,\n        array_aval=array_aval,\n        origin=origin,\n        pipeline_mode=self.pipeline_mode,\n        allow_captured_consts=allow_captured_consts,\n        debug=debug,\n    )\n    mapping.check_invariants()\n    return mapping\n","sourceCodeStart":680,"sourceCodeEnd":716,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/core.py#L680-L716","documentation":"Pallas index maps must be closed functions: their traced jaxpr must have no constants (closures). to_block_mapping raises when closed_jaxpr.consts is non-empty and allow_captured_consts is False (the default), because constants cannot always be handled during lowering/export.","triggerScenarios":"index_map closing over a Python/JAX value — e.g. block_shape known but index_map=lambda i: (i * block_size,) where block_size is captured from the enclosing scope instead of being passed/derived inside the map.","commonSituations":"Defining index maps inside factory functions that capture tile sizes or offsets; refactoring shared index helpers that close over state.","solutions":["Recompute needed values inside the index_map from its arguments instead of closing over them","Pass constants through as pallas_call inputs / out_shape plumbing if needed","If exporting and constants are intentional, use an API that sets allow_captured_consts (e.g. debug/export path)"],"exampleFix":"# before\ndef make_spec(bs):\n    return pl.BlockSpec((bs,), index_map=lambda i: (i * bs,))  # captures bs\n# after\ndef make_spec(bs):\n    return pl.BlockSpec((bs,), index_map=lambda i: (i * bs,) if False else (i * 128,))\n# better: derive from args: index_map=lambda i, bs=bs: (i * bs,)  # still captures; prefer explicit constants inline","handlingStrategy":"validation","validationCode":"import jax\nclosed = jax.make_jaxpr(index_map)(*grid_args)\nassert not closed.consts, f'index_map captures constants: {closed.consts}'","typeGuard":"def is_closed_index_map(index_map, args):\n    return not jax.make_jaxpr(index_map)(*args).consts","tryCatchPattern":null,"preventionTips":["Avoid closures in index maps; inline constants or derive from map args","Lint for lambda captures over mutable outer variables in kernel code"],"tags":["jax","pallas","index-map","closure"],"backgroundTag":"closure-captured-constant","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}