jax-ml/jax · error · ValueError

integer argument required; got dtype={arr.dtype}

Error message

integer argument required; got dtype={arr.dtype}

What it means

An internal JAX reductions helper requires an integer (or boolean) dtype array but received a floating-point or other non-integral dtype. This guard is used where indices/axes semantics are assumed.

Source

Thrown at jax/_src/numpy/reductions.py:215

                        else dtypes.finfo(a_dtype).max, dtype=a_dtype)
  try:
    return literals.TypedNdArray(np.array(init_val, dtype=a_dtype))
  except OverflowError:
    assert dtypes.issubdtype(a_dtype, np.integer)
    sign, info = np.sign(init_val), dtypes.iinfo(a_dtype)
    return literals.TypedNdArray(np.array(info.min if sign < 0 else info.max, dtype=a_dtype))

def _cast_to_bool(operand: Array) -> Array:
  if dtypes.issubdtype(operand.dtype, np.complexfloating):
    operand = operand.real
  return lax.convert_element_type(operand, np.bool_)

def _cast_to_numeric(operand: Array) -> Array:
  return promote_dtypes_numeric(operand)[0]

def _require_integer(arr: Array) -> Array:
  if not dtypes.isdtype(arr.dtype, ("bool", "integral")):
    raise ValueError(f"integer argument required; got dtype={arr.dtype}")
  return arr

def _ensure_optional_axes(x: Axis) -> Axis:
  def force(x):
    if x is None:
      return None
    try:
      return operator.index(x)
    except TypeError:
      return tuple(i if isinstance(i, str) else operator.index(i) for i in x)
  return core.concrete_or_error(
    force, x, "The axis argument must be known statically.")


@api.jit(static_argnames=('axis', 'dtype', 'keepdims', 'promote_integers'), inline=True)
def _reduce_sum(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
                out: None = None, keepdims: bool = False,
                initial: ArrayLike | None = None, where: ArrayLike | None = None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast the input to an integer dtype before the call: arr.astype(jnp.int32)
2. Check the argument types of the public API you called — this helper is internal, so the real mistake is in the caller's data
3. Inspect the stack trace to find which public jnp function routed here and validate its arguments

Example fix

// before
result = jnp.count_nonzero(mask, axis=axis)  # mask is float
// after
result = jnp.count_nonzero(mask.astype(bool), axis=axis)

arr = jnp.asarray(arr)
if not jnp.issubdtype(arr.dtype, jnp.integer):
    arr = arr.astype(jnp.int32)

def is_integer_array(a) -> bool:
    import jax.numpy as jnp
    return jnp.issubdtype(jnp.asarray(a).dtype, jnp.integer)

When it happens

Trigger: Triggered by internal paths in jax._src.numpy.reductions (e.g. argsort-adjacent or count-based helpers) when an array with float dtype reaches a helper expecting integral input; typically surfaced via public APIs like jnp.count_nonzero or reduction internals receiving float data where ints are required.

Common situations: Passing float arrays where an index/count array is expected; dtype promotion unexpectedly producing floats (e.g. via weak typing or jnp.asarray of Python floats).

Related errors

• jnp.{name}: where must be None or a boolean array; got {wher
• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• Accumulator aval mismatch: expected {aval}, got {acc.aval}
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f

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