{"record":{"id":"5825cf2fbc643f01","repo":"jax-ml/jax","slug":"reduction-axes-axes-contains-out-of-bounds-indic","errorCode":null,"errorMessage":"reduction axes {axes} contains out-of-bounds indices for {operand}.","messagePattern":"reduction axes (.+?) contains out-of-bounds indices for (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":8534,"sourceCode":"    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(\n          f\"{name} requires `out_sharding`'s unreduced_kind to be {out_kind}\"\n          f' but got {out_sharding.spec}')","sourceCodeStart":8516,"sourceCodeEnd":8552,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L8516-L8552","documentation":"A reduction primitive received an axis outside [0, operand.ndim). JAX checks every axis against the operand's rank before computing the output shape; both too-large positive and negative-beyond-rank indices fail.","triggerScenarios":"lax.reduce_max(x, axes=(2,)) on a 2-D array; jnp.sum(x, axis=3) where x.ndim==2; passing axis=-3 to a 2-D array via jnp.prod. Common when axis is computed from user input or config.","commonSituations":"Hardcoded axes that stop matching after adding/removing a batch dimension; loops that assume a fixed rank; passing a numpy-style axis that exceeds rank after jnp.squeeze/vmap transformations change the number of dimensions.","solutions":["Print x.ndim and the axes right before the call and clamp/validate: axes = tuple(a % x.ndim for a in axes).","If you removed a batch dim (squeeze/vmap), update hardcoded axis indices.","Validate user-supplied axis against the array rank: if not all(0 <= a < x.ndim for a in axes): raise ....","Use negative indexing intentionally (axis=-1 for last dim) instead of computing rank-dependent positives."],"exampleFix":"# before\nout = lax.reduce_sum(x, axes=(2,))  # x is 2-D\n# after\nout = lax.reduce_sum(x, axes=(-1,))  # last axis, rank-independent","handlingStrategy":"validation","validationCode":"assert all(0 <= a < x.ndim or -x.ndim <= a < 0 for a in axes), (x.shape, axes)\naxes = tuple(a % x.ndim for a in axes)\nout = lax.reduce_sum(x, axes)","typeGuard":"def axes_in_bounds(x, axes):\n    return all(-x.ndim <= a < x.ndim for a in axes)","tryCatchPattern":null,"preventionTips":["Prefer axis=-1/0-style relative indices over rank-dependent constants.","Log x.ndim alongside axes in debug builds."],"tags":["jax","lax","reduction","out-of-bounds","axis"],"backgroundTag":"axis-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}