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 a scalar

What it means

For Blocked/Element/Squeezed/int block dims, index_map must return scalar (shape-()) values. to_block_mapping checks each output aval: if it is not a ShapedArray at all (and has no shape), it raises this error at that position, since block indices must be scalars.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Return scalar values: use int(i) or .squeeze() / jnp scalar per dim
2. Match each output to its dim kind (Slice for BoundedSlice, scalar otherwise)
3. Print output avals if unsure by testing the map standalone under jax.make_jaxpr

Example fix

# before
index_map=lambda i: (i * 128 * jnp.ones(1),)
# after
index_map=lambda i: (i * 128,)

out = index_map(*idx)
flat = out if isinstance(out, tuple) else (out,)
assert all(not getattr(v, 'shape', ()) for v in flat), 'indices must be scalars'

def all_scalar_indices(vals):
    return all(getattr(v, 'ndim', 0) == 0 for v in vals)

When it happens

Trigger: index_map returning arrays/tensors (shape (1,) etc.), lists, or non-array values for scalar block dims — e.g. lambda i: (i[:, None],) or returning a Python tuple element.

Common situations: Vectorizing the index computation with extra brackets; returning jnp arrays of shape (1,) instead of scalars; mixing Slice and scalar positions.

Related errors

• Index map function {debug_info.func_src_info} for {origin} m
• index_map returned a value of type {type(idx_aval)} at posit
• Index map function {debug_info.func_src_info} for {origin} m
• Index map function {debug_info.func_src_info} for {origin} m
• Cannot swap scalars to VMEM.

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