jax-ml/jax · error · ValueError

`index_type` must be i32.

Error message

`index_type` must be i32.

What it means

Pallas arg-reduction lowering (argmax/argmin style ops) only supports int32 index outputs; the Triton index computation is hard-wired to i32. Passing index_dtype=int64/int16/uint32 raises ValueError before lowering.

Source

Thrown at jax/_src/pallas/triton/lowering.py:2540

  return _reduction_lowering(body, ctx, a, axes=axes)[0]


triton_lowering_rules[lax.reduce_max_p] = functools.partial(
    _reduce_lowering, jnp.maximum
)
triton_lowering_rules[lax.reduce_min_p] = functools.partial(
    _reduce_lowering, jnp.minimum
)
triton_lowering_rules[lax.reduce_sum_p] = functools.partial(
    _reduce_lowering, jnp.add
)


def _argreduce_lowering(
    body, ctx: LoweringRuleContext, a, *, axes, index_dtype
):
  if index_dtype != jnp.int32:
    raise ValueError("`index_type` must be i32.")
  if len(axes) != 1:
    raise ValueError("`pallas` reduce operations only support one reduce axis.")
  [axis] = axes
  [a_aval] = ctx.avals_in
  index = _make_range(0, a_aval.shape[axis])
  if len(a_aval.shape) > 1:
    # Broadcast index across the non-reduced axes
    for i in range(len(a_aval.shape)):
      if i != axis:
        index = _expand_dims(index, i)
    index = _bcast_to(index, a_aval.shape)
  ctx = ctx.replace(avals_in=[a_aval, a_aval.update(dtype=jnp.dtype(jnp.int32))])
  _, indices = _reduction_lowering(body, ctx, (a, index), axes=axes)
  return indices


def _reduce_argmax_combine(left, right):
  value1, index1 = left

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Omit index_dtype so the op uses int32, or set index_dtype=jnp.int32 explicitly
2. Cast any downstream index arithmetic back to the desired dtype after the arg-reduction

Example fix

# before
idx = jnp.argmax(x, index_dtype=jnp.int64)

# after
idx = jnp.argmax(x, index_dtype=jnp.int32)

import jax.numpy as jnp
assert index_dtype in (None, jnp.int32)

def valid_index_dtype(dt) -> bool:
    import jax.numpy as jnp
    return dt is None or dt == jnp.int32

When it happens

Trigger: Calling jnp.argmax(x, index_dtype=jnp.int64) or lax.argmax(..., index_dtype=...) inside a Mosaic triton kernel.

Common situations: Code written for XLA backends that accept int64 indices (default in some ops) then run under pallas; kernels indexing large arrays where int64 was chosen to avoid overflow assumptions.

Related errors

• Unsupported combination of input dtype ({x_aval.dtype}) and
• Only single axis reduction supported
• Only positional arguments are supported by debug_print on Pa
• unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}
• end must be greater than start, but got: {end} <= {start}

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