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 integer scalars. Output[{i}] has type {ov}.

What it means

Every value returned by a BlockSpec index_map must be an int32 or int64 scalar with empty shape. After tracing, to_block_mapping verifies each output aval's shape is empty and dtype is an integer type; otherwise it raises, naming the offending output index and its type.

Source

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

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

    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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Coerce to integer: use int step values or i * np.int32(step), and .astype(jnp.int32) on traced values
  2. Ensure returned avals are shape () scalars
  3. Keep all block-offset constants as Python ints

Example fix

# before
index_map=lambda b: (b * (128.0 / scale),)
# after
step = int(128 // scale)
index_map=lambda b: (b * step,)
Defensive patterns

Strategy: validation

Validate before calling

import jax, numpy as np
jaxpr = jax.make_jaxpr(index_map)(*map(jax.core.dim_constant, grid))
assert all(not oa.shape and oa.dtype in (np.int32, np.int64) for oa in jaxpr.out_avals)

Type guard

def integer_scalar_map(index_map, args):
    av = jax.make_jaxpr(index_map)(*args).out_avals
    return all(not a.shape and a.dtype in (jnp.int32, jnp.int64) for a in av)

Prevention

When it happens

Trigger: index_map returning floats (e.g. i * 0.5), booleans, or shaped arrays — e.g. lambda b: (b * step,) where step is a Python float, or returning a (1,) shaped array.

Common situations: Block-size arithmetic done with floats; index math producing bool from comparisons; returning indices computed in float32 for large-shape convenience.

Related errors

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