jax-ml/jax · error · ValueError

out_dtype must be integer typed; got {out_dtype=}

Error message

out_dtype must be integer typed; got {out_dtype=}

What it means

Raised by the Nonzero HiJAX primitive when the output dtype for the returned index arrays is not an integer type. Nonzero returns indices, so out_dtype must be integer-typed (typically int32).

Source

Thrown at jax/_src/numpy/hijax.py:203

  size: int
  axes: tuple[int, ...]
  out_dtype: np.dtype

  def __init__(
      self,
      a_aval: core.ShapedArray,
      *fill_value_avals: core.ShapedArray,
      size: int,
      axes: tuple[int, ...],
      out_dtype: np.dtype):
    if core.is_symbolic_dim(size):
      pass
    else:
      size = operator.index(size)
      if size < 0:
        raise ValueError(f"size must be a positive integer; got {size=}")
    if not dtypes.issubdtype(out_dtype, np.integer):
      raise ValueError(f"out_dtype must be integer typed; got {out_dtype=}")
    if not all(0 <= ax < a_aval.ndim for ax in axes):
      raise ValueError(f"axes out of range for array with {a_aval.ndim} dimensions:  {axes=}")
    if len(axes) != len(set(axes)):
      raise ValueError(f"duplicate axes are not allowed: {axes=}")
    if fill_value_avals and len(fill_value_avals) != len(axes):
      raise ValueError(f"Expected {len(axes)} fill values, got {len(fill_value_avals)}")
    if any(fv.dtype != out_dtype for fv in fill_value_avals):
      raise ValueError(f"Expected fill values to have dtype {out_dtype}, got {fill_value_avals}")
    batch_shape = tuple(
        s for i, s in enumerate(a_aval.shape) if i not in axes
    )
    for fv_aval in fill_value_avals:
      try:
        broadcasted = lax.broadcast_shapes(fv_aval.shape, batch_shape)
      except ValueError as e:
        raise ValueError(
            f"fill_value shape {fv_aval.shape} is not broadcast-compatible with "
            f"batch shape {batch_shape}"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use dtype='int32' or 'int64' (the default is int32)
2. If float output is needed, cast the returned index arrays after the call

Example fix

# before
idx = nonzero(a, size=10, dtype=jnp.float32)
# after
idx = nonzero(a, size=10, dtype=jnp.int32)

assert dtypes.issubdtype(np.dtype(dtype), np.integer), dtype

def is_int_dtype(d) -> bool:
    return dtypes.issubdtype(np.dtype(d), np.integer)

When it happens

Trigger: Calling nonzero(a, size=n, dtype='float32') or dtype=np.float64.

Common situations: Reusing a dtype config meant for data arrays; assuming dtype refers to the input array's dtype rather than the output indices.

Related errors

• Expected kind to be a dtype, string, or tuple; got {kind=}
• dtype must be specified.
• out_dtype should be an integer type; got {out_dtype}
• size must be a positive integer; got {size=}
• axes out of range for array with {a_aval.ndim} dimensions:

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