{"record":{"id":"2a1fb5646d395026","repo":"jax-ml/jax","slug":"duplicate-value-in-axes-of-reduction-axes","errorCode":null,"errorMessage":"duplicate value in 'axes' of reduction: {axes}","messagePattern":"duplicate value in 'axes' of reduction: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":8532,"sourceCode":"def _reduce_number_dtype_rule(name, operand, *_, **__):\n  if not dtypes.issubdtype(operand.dtype, np.number):\n    raise TypeError(\"{} does not accept dtype {}. Accepted dtypes are subtypes \"\n                    \"of number.\".format(name, dtype_to_string(operand.dtype)))\n  return operand.dtype\n\ndef _reduce_sum_transpose_rule(cotangent, operand, *, axes, out_sharding):\n  assert ad.is_undefined_primal(operand)\n  input_shape = operand.aval.shape\n  broadcast_dimensions = tuple(np.delete(np.arange(len(input_shape)), axes))\n  result = broadcast_in_dim(\n      cotangent, input_shape, broadcast_dimensions,\n      out_sharding=operand.aval.sharding)\n  assert result.shape == input_shape\n  return [result]\n\ndef _reduce_op_shape_rule(operand, *, axes, **_):\n  if len(axes) != len(set(axes)):\n    raise ValueError(f\"duplicate value in 'axes' of reduction: {axes}\")\n  if not all(0 <= a < operand.ndim for a in axes):\n    raise ValueError(f\"reduction axes {axes} contains out-of-bounds indices for {operand}.\")\n  axes = frozenset(axes)\n  return tuple(d for i, d in enumerate(operand.shape) if i not in axes)\n\ndef _reduce_op_sharding_rule_with_out_sharding(operand, *, axes, out_sharding):\n  if out_sharding is not None:\n    assert isinstance(out_sharding, NamedSharding)\n    return out_sharding\n  axes = frozenset(axes)\n  new_spec = P(*tuple(s for i, s in enumerate(operand.sharding.spec.partitions)\n                      if i not in axes))\n  return operand.sharding.update(spec=new_spec)\n\ndef _reduce_op_unreduced_rule(operand, axes, out_sharding, out_kind, name):\n  if out_sharding is not None and out_sharding.spec.unreduced:  # explicit mode\n    if out_sharding.spec.unreduced_kind is not out_kind:\n      raise core.ShardingTypeError(","sourceCodeStart":8514,"sourceCodeEnd":8550,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L8514-L8550","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Inspect the axes tuple passed to the reduction and remove duplicates (use tuple(set(axes)) or dict.fromkeys to preserve order).","If you built axes by concatenating lists, deduplicate before passing: axes=tuple(dict.fromkeys(a + b)).","Check any negative-axis normalization logic that may run twice and reintroduce a duplicate.","If the duplicate is unintended, review how axis lists are constructed upstream (loops, kwargs unpacking)."],"exampleFix":"// before\nout = lax.reduce_sum(x, axes=(1, 1))\n// after\nout = lax.reduce_sum(x, axes=(1,))","handlingStrategy":"validation","validationCode":"axes = tuple(dict.fromkeys(axes))  # dedupe, preserve order\nout = lax.reduce_sum(x, axes)","typeGuard":"def valid_axes(x, axes):\n    axes = tuple(dict.fromkeys(axes))\n    return all(0 <= a < x.ndim for a in axes)","tryCatchPattern":null,"preventionTips":["Never build axis lists by concatenation without deduplication.","Deduplicate axes with tuple(dict.fromkeys(...)) to keep deterministic order."],"tags":["jax","lax","reduction","axes","duplicate"],"backgroundTag":"invalid-axis-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}