jax-ml/jax · error · ValueError

duplicate axes are not allowed

Error message

duplicate axes are not allowed: {axes=}

What it means

The fill-value index helper treats the axes tuple as a set of distinct array dimensions to index. The constructor fires when the same axis value appears more than once in axes, which would select the same dimension twice and is not a valid indexing.

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
• Expected fill values, got
• fill_value must be a scalar or tuple of scalars; got
• fill_value tuple must have length equal to number of axes
• out_dtype must be integer typed; got

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

Appendix: 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)