jax-ml/jax · error · ValueError

BlockMapping for {self.origin} has captured constants: {self

Error message

BlockMapping for {self.origin} has captured constants: {self.index_map_jaxpr.consts}

What it means

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).

Source

Thrown at jax/_src/pallas/core.py:769

  See the `check_invariants` method for precise specification.
  """
  # TODO(apaszke,sharadmv): Replace mapped dims in block_shape with a transform.
  # After all, it's just indexing out singleton dimensions.
  block_shape: tuple[BlockDim, ...]
  transformed_block_aval: state.AbstractRef
  index_map_jaxpr: jax_core.Jaxpr
  index_map_out_tree: tree_util.PyTreeDef
  array_aval: jax_core.ShapedArray  # The whole array
  origin: OriginStr
  transforms: Sequence[state_types.Transform] = ()
  pipeline_mode: Buffered | None = None
  debug: bool = False
  allow_captured_consts: dataclasses.InitVar[bool] = False

  def __post_init__(self, allow_captured_consts: bool):
    if not allow_captured_consts and self.index_map_jaxpr.consts:
      raise ValueError(f"BlockMapping for {self.origin} has captured constants:"
                       f" {self.index_map_jaxpr.consts}")

  def check_invariants(self) -> None:
    if not config.enable_checks.value: return

    ref_block_shape = _get_ref_block_shape(self.block_shape)
    assert ref_block_shape == self.ref_aval.shape, (
        self.block_shape, self.ref_aval.shape)
    assert len(self.block_shape) == len(self.array_aval.shape), (
        self.block_shape, self.array_aval
    )
    assert all(ov.shape == () and
               (ov.dtype == jnp.int32 or ov.dtype == jnp.int64)
               for ov in self.index_map_jaxpr.out_avals), (
               self.index_map_jaxpr.out_avals)

  def replace(self, **kwargs):
    allow_captured_consts = len(self.index_map_jaxpr.consts) > 0

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Re-derive the BlockMapping via BlockSpec.to_block_mapping with allow_captured_consts set appropriately
2. Eliminate constants from the index map so the jaxpr is closed
3. Regenerate cached/exported artifacts with the current JAX version

mapping = spec.to_block_mapping(aval, ...)
assert not mapping.index_map_jaxpr.consts

When it happens

Trigger: 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.

Common situations: Version-skew between serialized Pallas programs and current JAX; tooling that reconstructs BlockMapping from cached jaxprs containing consts.

Related errors

• Unexpected dim_block_size: {dim_block_size}
• Unsupported block shape type: {type(dim)}
• Index map function {debug_info.func_src_info} for {origin} m
• The Pallas TPU lowering currently supports in memory space A
• Unsupported constant type: {x.type}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/383d24358b9bca7c. Report an issue: GitHub.