jax-ml/jax · error · TypeError

index_dtype must be an integer type, but got {}

Error message

index_dtype must be an integer type, but got {}

What it means

argmin/argmax's index_dtype parameter must be a NumPy integer type (e.g. int32/int64), but something else (float, bool, or a non-dtype) was supplied. This is a TypeError raised during dtype-rule evaluation.

Source

Thrown at jax/_src/lax/lax.py:8672

batching.defreducer(reduce_min_p)

def _argminmax_shape_rule(operand, *, axes, index_dtype):
  axis, = axes
  if not (0 <= axis < len(operand.shape)):
    raise ValueError(f"Invalid axis {axis} for operand shape {operand.shape}")
  if operand.shape[axis] < 1:
    raise ValueError("argmin and argmax require non-empty reduced dimension. "
                     f"operand.shape={operand.shape} {axis=}")
  return util.tuple_delete(operand.shape, axis)

def _argminmax_sharding_rule(operand, *, axes, index_dtype):
  axis, = axes
  return operand.sharding.update(spec=
      util.tuple_delete(operand.sharding.spec, axis))

def _argminmax_dtype_rule(operand, *, axes, index_dtype):
  if not dtypes.issubdtype(index_dtype, np.integer):
    raise TypeError("index_dtype must be an integer type, but got {}"
                    .format(dtype_to_string(index_dtype)))
  return index_dtype

class _ArgMinMaxReducer:

  def __init__(self, value_comparator: Callable[[Any, Any], Any]):
    self._value_comparator = value_comparator

  def __repr__(self):
    # Override the repr so that the metadata attached to the lowered op does not
    # contain unstable function ids. This plays more nicely with computation
    # fingerprint calculation in the compilation cache.
    return f'_ArgMinMaxReducer({self._value_comparator.__name__})'

  def __call__(self, op_val_index, acc_val_index):
    op_val, op_index = op_val_index
    acc_val, acc_index = acc_val_index
    # Pick op_val if Lt (for argmin) or if NaN

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a valid integer dtype string or object: index_dtype=jnp.int32 or 'int32'.
2. Check the value with dtypes.issubdtype(index_dtype, np.integer) if it comes from config.
3. Use jnp.argmin/jnp.argmax wrappers, which choose the index dtype automatically.

Example fix

# before
lax.argmin(x, 0, index_dtype=32)
# after
lax.argmin(x, 0, index_dtype=jnp.int32)

import numpy as np
from jax import dtypes
assert dtypes.issubdtype(index_dtype, np.integer), index_dtype
i = lax.argmin(x, 0, index_dtype=index_dtype)

def is_integer_dtype(dt):
    import numpy as np
    from jax import dtypes
    return dtypes.issubdtype(dt, np.integer)

When it happens

Trigger: lax.argmin(x, axis, index_dtype=jnp.float32), passing index_dtype='int' (not a real dtype), or passing a Python int like index_dtype=32 instead of a dtype object/string like 'int32'.

Common situations: Misreading the API and passing bit-width numbers or dtype names that don't exist; passing the platform default index type from another library.

Related errors

• Cannot interpret value of type {typ} as an abstract array; i
• {name} was requested to map a value of non-array type {core.
• primal and tangent arguments to jax.jvp must be tuples or li
• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass

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