jax-ml/jax · error · ValueError

axes out of range for array with {a_aval.ndim} dimensions:

Error message

axes out of range for array with {a_aval.ndim} dimensions:  {axes=}

What it means

Raised by the Nonzero HiJAX primitive when any axis in axes is negative or >= the input array's ndim. Note the nonzero() public function canonicalizes negative axes first, so this fires mainly when constructing the primitive directly with un-canonicalized axes.

Source

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

  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}"
        ) from e
      if broadcasted != batch_shape:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Validate axes against a.ndim: assert all(-a.ndim <= ax < a.ndim for ax in axes)
Use canonicalize_axis (jax._src.numpy.util.canonicalize_axis_tuple) to normalize negative axes before constructing the primitive

Example fix

# before
prim = Nonzero(aval, size=n, axes=(0, 2), out_dtype=np.int32)  # 2-D input
# after
from jax._src.numpy import util
axes = util.canonicalize_axis_tuple((0, 2), aval.ndim)  # choose valid axes
prim = Nonzero(aval, size=n, axes=axes, out_dtype=np.int32)

Defensive patterns

Strategy: validation

Validate before calling

assert all(0 <= ax < a.ndim for ax in axes), (a.ndim, axes)

Type guard

def axes_in_range(axes, ndim: int) -> bool:
    return all(0 <= ax < ndim for ax in axes)

Prevention

Use util.canonicalize_axis_tuple to normalize negative axes
Assert rank assumptions near tensor-constructing code

When it happens

Trigger: Constructing Nonzero(a_aval, axes=(0, 2), ...) for a 2-D array; passing axes larger than ndim-1; axes for an input whose ndim shrank after a refactor.

Common situations: Hardcoded axes that outlived a shape change; passing axes meant for a different rank of tensor in generic library code.

Related errors

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