jax-ml/jax · error · ValueError

duplicate value in 'axes' of reduction: {axes}

Error message

duplicate value in 'axes' of reduction: {axes}

What it means

Raised by JAX's lax reduction shape rule when the 'axes' argument to a reduction primitive (reduce_sum, reduce_max, etc.) contains the same axis more than once. JAX requires axes to be a set of distinct integers; duplicates are ambiguous and rejected before shape inference.

Source

Thrown at jax/_src/lax/lax.py:8532

def _reduce_number_dtype_rule(name, operand, *_, **__):
  if not dtypes.issubdtype(operand.dtype, np.number):
    raise TypeError("{} does not accept dtype {}. Accepted dtypes are subtypes "
                    "of number.".format(name, dtype_to_string(operand.dtype)))
  return operand.dtype

def _reduce_sum_transpose_rule(cotangent, operand, *, axes, out_sharding):
  assert ad.is_undefined_primal(operand)
  input_shape = operand.aval.shape
  broadcast_dimensions = tuple(np.delete(np.arange(len(input_shape)), axes))
  result = broadcast_in_dim(
      cotangent, input_shape, broadcast_dimensions,
      out_sharding=operand.aval.sharding)
  assert result.shape == input_shape
  return [result]

def _reduce_op_shape_rule(operand, *, axes, **_):
  if len(axes) != len(set(axes)):
    raise ValueError(f"duplicate value in 'axes' of reduction: {axes}")
  if not all(0 <= a < operand.ndim for a in axes):
    raise ValueError(f"reduction axes {axes} contains out-of-bounds indices for {operand}.")
  axes = frozenset(axes)
  return tuple(d for i, d in enumerate(operand.shape) if i not in axes)

def _reduce_op_sharding_rule_with_out_sharding(operand, *, axes, out_sharding):
  if out_sharding is not None:
    assert isinstance(out_sharding, NamedSharding)
    return out_sharding
  axes = frozenset(axes)
  new_spec = P(*tuple(s for i, s in enumerate(operand.sharding.spec.partitions)
                      if i not in axes))
  return operand.sharding.update(spec=new_spec)

def _reduce_op_unreduced_rule(operand, axes, out_sharding, out_kind, name):
  if out_sharding is not None and out_sharding.spec.unreduced:  # explicit mode
    if out_sharding.spec.unreduced_kind is not out_kind:
      raise core.ShardingTypeError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Inspect the axes tuple passed to the reduction and remove duplicates (use tuple(set(axes)) or dict.fromkeys to preserve order).
2. If you built axes by concatenating lists, deduplicate before passing: axes=tuple(dict.fromkeys(a + b)).
3. Check any negative-axis normalization logic that may run twice and reintroduce a duplicate.
4. If the duplicate is unintended, review how axis lists are constructed upstream (loops, kwargs unpacking).

Example fix

// before
out = lax.reduce_sum(x, axes=(1, 1))
// after
out = lax.reduce_sum(x, axes=(1,))

axes = tuple(dict.fromkeys(axes))  # dedupe, preserve order
out = lax.reduce_sum(x, axes)

def valid_axes(x, axes):
    axes = tuple(dict.fromkeys(axes))
    return all(0 <= a < x.ndim for a in axes)

When it happens

Trigger: Calling lax.reduce_sum(x, axes=(1,1)), jnp.sum(x, axis=(0,0)) with a repeated entry, or any reduce window/argop with duplicate axis values. Also occurs when axes are built dynamically (e.g. tuple(range(n)) + (0,)) and accidentally overlap.

Common situations: Programmatically composing axis lists (concatenating per-term axes), negative-axis normalization applied twice, or copy-pasted axis tuples. Newer JAX versions validate this eagerly during tracing rather than at compile time, surfacing errors earlier.

Related errors

• reduction axes {axes} contains out-of-bounds indices for {op
• reductions require axes to be (0,) on SparseCore, but got {a
• `pallas` reduce operations only support one reduce axis.
• lax.platform_dependent: the '{pname}' branch must be a calla
• Use 'cuda', 'rocm', or 'oneapi' for lax.platform_dependent.

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