{"record":{"id":"383d24358b9bca7c","repo":"jax-ml/jax","slug":"blockmapping-for-self-origin-has-captured-consta","errorCode":null,"errorMessage":"BlockMapping for {self.origin} has captured constants: {self.index_map_jaxpr.consts}","messagePattern":"BlockMapping for (.+?) has captured constants: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/core.py","lineNumber":769,"sourceCode":"\n  See the `check_invariants` method for precise specification.\n  \"\"\"\n  # TODO(apaszke,sharadmv): Replace mapped dims in block_shape with a transform.\n  # After all, it's just indexing out singleton dimensions.\n  block_shape: tuple[BlockDim, ...]\n  transformed_block_aval: state.AbstractRef\n  index_map_jaxpr: jax_core.Jaxpr\n  index_map_out_tree: tree_util.PyTreeDef\n  array_aval: jax_core.ShapedArray  # The whole array\n  origin: OriginStr\n  transforms: Sequence[state_types.Transform] = ()\n  pipeline_mode: Buffered | None = None\n  debug: bool = False\n  allow_captured_consts: dataclasses.InitVar[bool] = False\n\n  def __post_init__(self, allow_captured_consts: bool):\n    if not allow_captured_consts and self.index_map_jaxpr.consts:\n      raise ValueError(f\"BlockMapping for {self.origin} has captured constants:\"\n                       f\" {self.index_map_jaxpr.consts}\")\n\n  def check_invariants(self) -> None:\n    if not config.enable_checks.value: return\n\n    ref_block_shape = _get_ref_block_shape(self.block_shape)\n    assert ref_block_shape == self.ref_aval.shape, (\n        self.block_shape, self.ref_aval.shape)\n    assert len(self.block_shape) == len(self.array_aval.shape), (\n        self.block_shape, self.array_aval\n    )\n    assert all(ov.shape == () and\n               (ov.dtype == jnp.int32 or ov.dtype == jnp.int64)\n               for ov in self.index_map_jaxpr.out_avals), (\n               self.index_map_jaxpr.out_avals)\n\n  def replace(self, **kwargs):\n    allow_captured_consts = len(self.index_map_jaxpr.consts) > 0","sourceCodeStart":751,"sourceCodeEnd":787,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/core.py#L751-L787","documentation":"BlockMapping is the internal record produced by BlockSpec.to_block_mapping. Its __post_init__ enforces that index_map_jaxpr has no captured constants unless allow_captured_consts was explicitly set, mirroring the export-time check but at dataclass construction (e.g. deserialization or manual construction).","triggerScenarios":"Directly constructing or (de)serializing a BlockMapping whose index_map_jaxpr.consts is non-empty without allow_captured_consts=True — typically in export/import pipelines or code that rebuilds BlockMappings.","commonSituations":"Version-skew between serialized Pallas programs and current JAX; tooling that reconstructs BlockMapping from cached jaxprs containing consts.","solutions":["Re-derive the BlockMapping via BlockSpec.to_block_mapping with allow_captured_consts set appropriately","Eliminate constants from the index map so the jaxpr is closed","Regenerate cached/exported artifacts with the current JAX version"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"mapping = spec.to_block_mapping(aval, ...)\nassert not mapping.index_map_jaxpr.consts","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Construct BlockMappings only via BlockSpec.to_block_mapping","Regenerate serialized Pallas artifacts on JAX upgrades"],"tags":["jax","pallas","block-mapping","internal","closure"],"backgroundTag":"closure-captured-constant","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}