jax-ml/jax · error · ValueError

duplicate axes are not allowed: {axes=}

Error message

duplicate axes are not allowed: {axes=}

What it means

Raised by the Nonzero HiJAX primitive when the axes tuple contains duplicate entries. Each axis in the reduction set must be unique because each axis produces exactly one output index array.

Source

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

  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:
        raise ValueError(
            f"fill_value shape {fv_aval.shape} cannot be broadcast to "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Deduplicate while preserving order: axes = tuple(dict.fromkeys(axes))
2. Review the code that builds the axes tuple for accidental repetition

Example fix

# before
axes = (axis, axis)
# after
axes = (axis,)
# or dedupe: axes = tuple(dict.fromkeys(axes))

assert len(axes) == len(set(axes)), axes

def unique_axes(axes) -> bool:
    return len(axes) == len(set(axes))

When it happens

Trigger: Constructing Nonzero(..., axes=(1, 1)) or building axes dynamically so the same axis appears twice, e.g. axes = (axis, axis).

Common situations: Programmatically composing axes lists that concatenate shared axes without deduplication.

Related errors

• axes out of range for array with {a_aval.ndim} dimensions:
• size must be a positive integer; got {size=}
• out_dtype must be integer typed; got {out_dtype=}
• Expected {len(axes)} fill values, got {len(fill_value_avals)
• fill_value tuple must have length equal to number of axes ({

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