jax-ml/jax · error · ValueError

index_map returned a value of type {type(idx_aval)} at posit

Error message

index_map returned a value of type {type(idx_aval)} at position {i} with block dimension {bd} when it should be pl.Slice

What it means

For block dims declared as BoundedSlice, the index_map must return a pl.Slice (jax_core indexing.Slice) object at that position, because a bounded slice needs start/size info. Returning an integer or anything else for a BoundedSlice dim triggers this error.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Return jax.lax.Slice or pl.slice objects for BoundedSlice dims as the docs/examples show
2. Return ints only for Blocked/Element/Squeezed dims
3. Mirror the canonical BoundedSlice example from the Pallas test suite

Example fix

# before
spec = pl.BlockSpec(block_shape=(pl.BoundedSlice(128),),
                    index_map=lambda i: (i,))
# after
spec = pl.BlockSpec(block_shape=(pl.BoundedSlice(128),),
                    index_map=lambda i: (pl.slice(i * 128, size=128),))

import jax.experimental.pallas as pl
out = index_map(*indices)
for v, bd in zip(out if isinstance(out, tuple) else (out,), block_shape):
    if isinstance(bd, pl.BoundedSlice):
        assert isinstance(v, jax.lax.Slice), 'BoundedSlice dim needs pl.Slice'

def slice_for_bounded(block_shape):
    return any(isinstance(d, pl.BoundedSlice) for d in block_shape)

When it happens

Trigger: Using pl.BoundedSlice in block_shape but returning plain ints from index_map, or returning a Slice where an int is expected on other dims (this specific branch covers BoundedSlice dims).

Common situations: Adopting BoundedSlice for variable-length blocks (e.g. ragged/attention kernels) without updating the index map to produce pl.Slice instances.

Related errors

• Index map function {debug_info.func_src_info} for {origin} m
• Unsupported block dimension type: {type(dim)}. Allowed types
• Unsupported block shape type: {type(dim)}
• Array type must have a `shape` attribute, but got {type(arra
• Block shape for {origin} (= {block_shape}) must have the sam

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