jax-ml/jax · error · ValueError

Index map function {debug_info.func_src_info} for {origin} m

Error message

Index map function {debug_info.func_src_info} for {origin} must not capture constants: {closed_jaxpr.consts}

What it means

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.

Source

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

          if (
              not isinstance(idx_aval, jax_core.ShapedArray)
              and not idx_aval.shape
          ):
            raise ValueError(
                "index_map returned a value of type"
                f" {type(idx_aval)} at position {i} with block dimension"
                f" {bd} when it should be a scalar"
            )
    for i, ov in enumerate(out_avals):
      if ov.shape or ov.dtype not in [jnp.int32, jnp.int64]:
        raise ValueError(
            f"Index map function {debug_info.func_src_info} for "
            f"{origin} must return integer scalars. Output[{i}] has type "
            f"{ov}."
        )

    if closed_jaxpr.consts and not allow_captured_consts:
      raise ValueError(
          f"Index map function {debug_info.func_src_info} for "
          f"{origin} must not capture constants: {closed_jaxpr.consts}"
      )

    mapping = BlockMapping(
        block_shape=block_shape,
        transformed_block_aval=block_aval,  # There are no transforms by default
        index_map_jaxpr=closed_jaxpr,
        index_map_out_tree=out_avals.tree,
        array_aval=array_aval,
        origin=origin,
        pipeline_mode=self.pipeline_mode,
        allow_captured_consts=allow_captured_consts,
        debug=debug,
    )
    mapping.check_invariants()
    return mapping

View on GitHub (pinned to 1e1c6a8fc0)

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)

Example fix

# before
def make_spec(bs):
    return pl.BlockSpec((bs,), index_map=lambda i: (i * bs,))  # captures bs
# after
def make_spec(bs):
    return pl.BlockSpec((bs,), index_map=lambda i: (i * bs,) if False else (i * 128,))
# better: derive from args: index_map=lambda i, bs=bs: (i * bs,)  # still captures; prefer explicit constants inline

Defensive patterns

Strategy: validation

Validate before calling

import jax
closed = jax.make_jaxpr(index_map)(*grid_args)
assert not closed.consts, f'index_map captures constants: {closed.consts}'

Type guard

def is_closed_index_map(index_map, args):
    return not jax.make_jaxpr(index_map)(*args).consts

Prevention

Avoid closures in index maps; inline constants or derive from map args
Lint for lambda captures over mutable outer variables in kernel code

When it happens

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

Common situations: Defining index maps inside factory functions that capture tile sizes or offsets; refactoring shared index helpers that close over state.

Related errors

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