Repository: google/jax

Issue:
jax.numpy.sum behaving funny for non-existent axis
When invoking the `sum` operation on a non-existent axis I don't get an error (as would be the expected result), but sums over one of the existing axes. As far as I can tell, for any array `x` with `n` axes, calls to `sum` with axis `j > n` behave as if the provided axis was `j % n`. E.g., if the array `x` has 3 axes, calls to `sum` with `j` being any multiple of three all sum over the first axis (j=0,3,6,9,... all give the same result).
This is especially noteworthy for the one-dimensional array (only one axis), where calls to `sum` with any value for `axis` always sums over the whole array.
It would be preferable if either an error was raised, or summing over a non-existing axis would be a no-op.

Might be related to #278 

example code
```
import jax.numpy as np
x = np.array([ [ [ 1, 2 ], [ 3, 4 ] ], [ [ 5, 6 ], [ 7, 8] ], [ [ 9, 10 ], [ 11, 12 ] ] ])
assert(np.all(np.sum(x, axis=3) == np.sum(x, axis=0)))
assert(np.all(np.sum(x, axis=4) == np.sum(x, axis=1)))
assert(np.all(np.sum(x, axis=5) == np.sum(x, axis=2)))
assert(np.all(np.sum(x, axis=6) == np.sum(x, axis=0)))
# and so on..
```


Discussion:

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -231,6 +231,17 @@ def wrap(op):
       return op
   return wrap
 
+# TODO(phawkins): use this helper everywhere.
+def _canonicalize_axis(axis, num_dims):
+  """Canonicalize an axis in (-num_dims, num_dims) to [0, num_dims)."""
+  axis = int(axis)
+  if axis < 0:
+    axis = axis + num_dims
+  if axis < 0 or axis >= num_dims:
+      raise ValueError(
+          "axis {} is out of bounds for array of dimension {}".format(
+              axis, num_dims))
+  return axis
 
 ### implementations of numpy functions in terms of lax
 
@@ -944,9 +955,9 @@ def _reduction_dims(a, axis):
   if axis is None:
     return onp.arange(ndim(a))
   elif isinstance(axis, (onp.ndarray, tuple, list)):
-    return onp.mod(onp.asarray(axis), ndim(a))
+    return tuple(_canonicalize_axis(x, ndim(a)) for x in axis)
   elif isinstance(axis, int):
-    return onp.mod([axis], ndim(a))
+    return (_canonicalize_axis(axis, ndim(a)),)
   else:
     raise TypeError("Unexpected type of axis argument: {}".format(type(axis)))
 
@@ -1154,19 +1165,6 @@ def cumulative_reduction(a, axis=None, dtype=None):
 
 ### Array-creation functions
 
-# TODO(phawkins): use this helper everywhere.
-def _canonicalize_axis(axis, num_dims):
-  """Canonicalize an axis in (-num_dims, num_dims) to [0, num_dims)."""
-  axis = int(axis)
-  if axis < 0:
-    axis = axis + num_dims
-  if axis < 0 or axis >= num_dims:
-      raise ValueError(
-          "axis {} is out of bounds for array of dimension {}".format(
-              axis, num_dims))
-  return axis
-
-
 @_wraps(onp.pad)
 def pad(array, pad_width, mode, constant_values=0):
   if mode != "constant":


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1622,6 +1622,10 @@ def f(x, v):
     first_call = f(x, v)
     second_call = f(x, v)  # doesn't crash
 
+  def testReductionOfOutOfBoundsAxis(self):  # Issue 888
+    x = lnp.ones((3, 4))
+    self.assertRaises(ValueError, lambda: lnp.sum(x, axis=2))
+
 
 if __name__ == "__main__":
   absltest.main()