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 return {len(block_shape)} values to match {block_shape=}. Currently returning {len(unflat_avals)} values:

What it means

The BlockSpec index_map must return exactly one index value per block dimension. to_block_mapping traces the index_map, flattens its outputs, and compares the count with len(block_shape); mismatch raises this error including the function's source info and both counts.

Source

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

      )

    fake_index_map_args, fake_index_map_kwargs = \
        index_map_tree.unflatten([False] * index_map_tree.num_leaves)
    debug_info = api_util.debug_info(
        "pallas_call index_map",
        index_map_func,
        fake_index_map_args,
        fake_index_map_kwargs,
    )
    with tracing_grid_env(grid, vmapped_dims):
      closed_jaxpr, out_avals = pe.trace_to_jaxpr(
          index_map_func,
          ft.FTPyTree(index_map_avals, index_map_tree),
          debug_info)
    unflat_avals = out_avals.unflatten()

    if len(unflat_avals) != len(block_shape):
      raise ValueError(
          f"Index map function {debug_info.func_src_info} for "
          f"{origin} must return "
          f"{len(block_shape)} values to match {block_shape=}. "
          f"Currently returning {len(unflat_avals)} values:"
      )
    # Verify types match
    for i, (idx_aval, bd) in enumerate(zip(unflat_avals, block_shape)):
      match bd:
        case BoundedSlice():
          if not isinstance(idx_aval, indexing.Slice):
            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 pl.Slice"
            )
        case Blocked() | Element() | Squeezed() | int():
          if (
              not isinstance(idx_aval, jax_core.ShapedArray)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Return exactly len(block_shape) values, one per dim (use 0 for Squeezed dims)
Remove nested parentheses: return i, j not ((i, j),)
Check the counts in the error message to see which operand's map is wrong

Example fix

# before
pl.BlockSpec(block_shape=(None, 128), index_map=lambda i: (i, 0, 0))
# after
pl.BlockSpec(block_shape=(None, 128), index_map=lambda i: (i, 0))

Defensive patterns

Strategy: validation

Validate before calling

import jax
jaxpr = jax.make_jaxpr(index_map)(*grid_indices)
n_out = len(jaxpr.out_avals)
assert n_out == len(block_shape), f'index_map returns {n_out}, block has {len(block_shape)}'

Prevention

Return flat tuples: return i, j — never ((i, j),)
Count placeholders (0) for Squeezed dims when writing maps

When it happens

Trigger: index_map returning a single tuple (lambda i, j: ((i, j),)) vs multiple values, returning the wrong number of elements, or returning None/a scalar for a multi-dim block shape (e.g. block_shape=(None, 128) with index_map=lambda i: (i, 0)).

Common situations: Extra parentheses around tuple returns; forgetting the squeezed dims still need a placeholder index; refactoring kernels from 1D to 2D without updating the map's return arity.

Related errors

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