jax-ml/jax · error · ValueError

duplicate value in 'axis': {axis}

Error message

duplicate value in 'axis': {axis}

What it means

The axis argument to a reduction contained the same axis more than once after canonicalization (including negative indices resolving to the same axis, or duplicated names), which numpy also rejects.

Source

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

    result = op(initial_arr, result)
  if keepdims:
    result = lax.expand_dims(result, pos_dims)
  return lax.convert_element_type(result, dtype or result_dtype)

def _canonicalize_axis_allow_named(x, rank):
  return maybe_named_axis(x, lambda i: canonicalize_axis(i, rank), lambda name: name)

def _reduction_dims(a: ArrayLike, axis: Axis):
  if axis is None:
    return (tuple(range(np.ndim(a))),) * 2
  if not isinstance(axis, (np.ndarray, tuple, list)):
    axes = (axis,)
  else:
    axes = axis
  canon_axis = tuple(_canonicalize_axis_allow_named(x, np.ndim(a))
                     for x in axes)
  if len(canon_axis) != len(set(canon_axis)):
    raise ValueError(f"duplicate value in 'axis': {axis}")
  canon_pos_axis = tuple(x for x in canon_axis if isinstance(x, int))
  if len(canon_pos_axis) != len(canon_axis):
    return canon_pos_axis, canon_axis
  else:
    return canon_axis, canon_axis

def _reduction_init_val(a: Array, init_val: Any) -> np.ndarray:
  # This function uses np.* functions because lax pattern matches against the
  # specific concrete values of the reduction inputs. TypedNdArray prevents
  # canonicalization when explicit 64-bit dtypes are allowed.
  a_dtype = a.dtype
  if a_dtype == 'bool':
    return literals.TypedNdArray(np.array(init_val > 0, dtype=a_dtype))
  if (np.isinf(init_val) and dtypes.issubdtype(a_dtype, np.floating)
      and not dtypes.supports_inf(a_dtype)):
    init_val = np.array(dtypes.finfo(a_dtype).min if np.isneginf(init_val)
                        else dtypes.finfo(a_dtype).max, dtype=a_dtype)
  try:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Deduplicate axes before calling: axis=tuple(set(canonical_axes))
  2. Normalize negative indices first, then dedupe
  3. Validate the axis tuple with a helper before the reduction

Example fix

// before
jnp.sum(x, axis=(0, 0, 1))
// after
axes = tuple(dict.fromkeys(ax % x.ndim for ax in (0, 0, 1)))
jnp.sum(x, axis=axes)
Defensive patterns

Strategy: validation

Validate before calling

axes = tuple(dict.fromkeys(a % x.ndim for a in axis_tuple))  # dedupe after normalizing negatives
jnp.sum(x, axis=axes)

Prevention

When it happens

Trigger: jnp.sum(x, axis=(0, 0)), jnp.sum(x, axis=(0, -2)) on a 2-D+ array, or axis=('batch', 0) where 'batch' canonicalizes to 0.

Common situations: Programmatically building axis tuples that concatenate ranges without deduplication; handling negative indices alongside positive ones (axis=(1, -1) on 2-D input).

Related errors


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