Repository: google/jax
Issue: jax.numpy.average can't take tuple axis

### Description

Tuple axis is documented: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.average.html#jax.numpy.average

And supported in numpy:

```python
import numpy as np
np.average(np.arange(0, 4), axis=(0,), weights=np.arange(1, 5))  # 2.0
```

But will result an error

```python
from jax import numpy as jnp
jnp.average(jnp.arange(0, 4), axis=(0,), weights=jnp.arange(1, 5))
```

```
TypeError: 'tuple' object cannot be interpreted as an integer
```

Possibly related: #6451

### What jax/jaxlib version are you using?

jax v0.3.17, jaxlib v0.3.15

### Which accelerator(s) are you using?

CPU

### Additional System Info

Linux

Changed files:
jax/_src/test_util.py, tests/lax_numpy_operators_test.py, tests/lax_numpy_test.py, jax/_src/numpy/reductions.py, tests/lax_numpy_reducers_test.py

diff --git a/jax/_src/numpy/reductions.py b/jax/_src/numpy/reductions.py
index e3b090154..99a22b62f 100644
--- a/jax/_src/numpy/reductions.py
+++ b/jax/_src/numpy/reductions.py
@@ -20,7 +20,6 @@ import warnings
 
 import numpy as np
 
-import jax
 from jax import core
 from jax import lax
 from jax._src import api
@@ -28,7 +27,7 @@ from jax._src import dtypes
 from jax._src.numpy.ndarray import ndarray
 from jax._src.numpy.util import _broadcast_to, _check_arraylike, _complex_elem_type, _promote_dtypes_inexact, _promote_dtypes_numeric, _where, _wraps
 from jax._src.lax import lax as lax_internal
-from jax._src.util import canonicalize_axis as _canonicalize_axis, maybe_named_axis
+from jax._src.util import canonicalize_axis as _canonicalize_axis, maybe_named_axis, prod as _prod
 
 
 _all = builtins.all
@@ -328,8 +327,10 @@ def _average(a, axis: Optional[Union[int, Tuple[int, ...]]] = None, weights=None
     avg = mean(a, axis=axis, keepdims=keepdims)
     if axis is None:
       weights_sum = lax.full((), core.dimension_as_value(a.size), dtype=avg.dtype)
+    elif isinstance(axis, tuple):
+      weights_sum = lax.full_like(avg, _prod(core.dimension_as_value(a.shape[d]) for d in axis))
     else:
-      weights_sum = lax.full_like(avg, core.dimension_as_value(a.shape[axis]), dtype=avg.dtype)
+      weights_sum = lax.full_like(avg, core.dimension_as_value(a.shape[axis]))
   else:
     _check_arraylike("average", a, weights)
     a, weights = _promote_dtypes_inexact(a, weights)
@@ -337,17 +338,25 @@ def _average(a, axis: Optional[Union[int, Tuple[int, ...]]] = None, weights=None
     a_shape = np.shape(a)
     a_ndim = len(a_shape)
     weights_shape = np.shape(weights)
-    axis = None if axis is None else _canonicalize_axis(axis, a_ndim)
+
+    if axis is None:
+      pass
+    elif isinstance(axis, tuple):
+      axis = tuple(_canonicalize_axis(d, a_ndim) for d in axis)
+    else:
+      axis = _canonicalize_axis(axis, a_ndim)
 
     if a_shape != weights_shape:
       # Make sure the dimensions work out
-      if axis is None:
-        raise ValueError("Axis must be specified when shapes of a and "
-                         "weights differ.")
       if len(weights_shape) != 1:
         raise ValueError("1D weights expected when shapes of a and "
                          "weights differ.")
-      if not core.symbolic_equal_dim(weights_shape[0], a_shape[axis]):
+      if axis is None:
+        raise ValueError("Axis must be specified when shapes of a and "
+                         "weights differ.")
+      elif isinstance(axis, tuple):
+        raise ValueError("Single axis expected when shapes of a and weights differ")
+      elif not core.symbolic_equal_dim(weights_shape[0], a_shape[axis]):
         raise ValueError("Length of weights not "
                          "compatible with specified axis.")
 
diff --git a/jax/_src/test_util.py b/jax/_src/test_util.py
index 769bffe73..2c3aae52c 100644
--- a/jax/_src/test_util.py
+++ b/jax/_src/test_util.py
@@ -37,8 +37,9 @@ from jax import core
 from jax._src import dtypes as _dtypes
 from jax import lax
 from jax._src.config import flags, bool_env, config
+from jax._src.numpy.lax_numpy import _promote_dtypes, _promote_dtypes_inexact
 from jax._src.util import prod, unzip2
-from jax.tree_util import tree_map, tree_all
+from jax.tree_util import tree_map, tree_all, tree_flatten, tree_unflatten
 from jax._src.lib import xla_bridge
 from jax._src import dispatch
 from jax._src.public_test_util import (  # noqa: F401
@@ -758,6 +759,21 @@ def with_config(**kwds):
   return decorator
 
 
+def promote_like_jnp(fun, inexact=False):
+  """Decorator that promotes the arguments of `fun` to `jnp.result_type(*args)`.
+
+  jnp and np have different type promotion semantics; this decorator allows
+  tests make an np reference implementation act more like an jnp
+  implementation.
+  """
+  _promote = _promote_dtypes_inexact if inexact else _promote_dtypes
+  def wrapper(*args, **kw):
+    flat_args, tree = tree_flatten(args)
+    args = tree_unflatten(tree, _promote(*flat_args))
+    return fun(*args, **kw)
+  return wrapper
+
+
 class JaxTestCase(parameterized.TestCase):
   """Base class for JAX tests including numerical checks and boilerplate."""
   _default_config = {
diff --git a/tests/lax_numpy_operators_test.py b/tests/lax_numpy_operators_test.py
index 0b3d4e4db..4496ef023 100644
--- a/tests/lax_numpy_operators_test.py
+++ b/tests/lax_numpy_operators_test.py
@@ -29,11 +29,9 @@ import jax
 import jax.ops
 from jax import lax
 from jax import numpy as jnp
-from jax import tree_util
 
 from jax._src import dtypes
 from jax._src import test_util as jtu
-from jax._src.numpy.lax_numpy import _promote_dtypes, _promote_dtypes_inexact
 
 from jax.config import config
 config.parse_flags_with_absl()
@@ -66,13 +64,6 @@ python_scalar_dtypes = [jnp.bool_, jnp.int_, jnp.float_, jnp.complex_]
 # uint64 is problematic because with any uint type it promotes to float:
 int_dtypes_no_uint64 = [d for d in int_dtypes + unsigned_dtypes if d != np.uint64]
 
-def _indexer_with_default_outputs(indexer, use_defaults=True):
-  """Like jtu.with_jax_dtype_defaults, but for __getitem__ APIs"""
-  class Indexer:
-    @partial(jtu.with_jax_dtype_defaults, use_defaults=use_defaults)
-    def __getitem__(self, *args):
-      return indexer.__getitem__(*args)
-  return Indexer()
 
 def _valid_dtypes_for_shape(shape, dtypes):
   # Not all (shape, dtype) pairs are valid. In particular, Python scalars only
@@ -400,21 +391,6 @@ def _shapes_are_equal_length(shapes):
   return all(len(shape) == len(shapes[0]) for shape in shapes[1:])
 
 
-def _promote_like_jnp(fun, inexact=False):
-  """Decorator that promotes the arguments of `fun` to `jnp.result_type(*args)`.
-
-  jnp and np have different type promotion semantics; this decorator allows
-  tests make an np reference implementation act more like an jnp
-  implementation.
-  """
-  _promote = _promote_dtypes_inexact if inexact else _promote_dtypes
-  def wrapper(*args, **kw):
-    flat_args, tree = tree_util.tree_flatten(args)
-    args = tree_util.tree_unflatten(tree, _promote(*flat_args))
-    return fun(*args, **kw)
-  return wrapper
-
-
 class JaxNumpyOperatorTests(jtu.JaxTestCase):
   """Tests for LAX-backed Numpy operators."""
 
@@ -460,7 +436,7 @@ class JaxNumpyOperatorTests(jtu.JaxTestCase):
                            [tolerance, tol, jtu.default_tolerance()])
 
     with jtu.strict_promotion_if_dtypes_match(dtypes):
-      self._CheckAgainstNumpy(_promote_like_jnp(np_op, inexact), jnp_op,
+      self._CheckAgainstNumpy(jtu.promote_like_jnp(np_op, inexact), jnp_op,
                               args_maker, check_dtypes=check_dtypes, tol=tol)
       self._CompileAndCheck(jnp_op, args_maker, check_dtypes=check_dtypes,
                             atol=tol, rtol=tol)
@@ -595,7 +571,7 @@ class JaxNumpyOperatorTests(jtu.JaxTestCase):
     rng = rng_factory(self.rng())
     args_maker = self._GetArgsMaker(rng, shapes, dtypes)
     with jtu.strict_promotion_if_dtypes_match(dtypes):
-      self._CheckAgainstNumpy(_promote_like_jnp(np_op), jnp_op, args_maker)
+      self._CheckAgainstNumpy(jtu.promote_like_jnp(np_op), jnp_op, args_maker)
       self._CompileAndCheck(jnp_op, args_maker)
 
   @parameterized.named_parameters(jtu.cases_from_list(
diff --git a/tests/lax_numpy_reducers_test.py b/tests/lax_numpy_reducers_test.py
index 6444b81e4..3f70384f8 100644
--- a/tests/lax_numpy_reducers_test.py
+++ b/tests/lax_numpy_reducers_test.py
@@ -57,6 +57,18 @@ all_dtypes = number_dtypes + bool_dtypes
 
 python_scalar_dtypes = [jnp.bool_, jnp.int_, jnp.float_, jnp.complex_]
 
+def _valid_dtypes_for_shape(shape, dtypes):
+  # Not all (shape, dtype) pairs are valid. In particular, Python scalars only
+  # have one type in each category (float, bool, etc.)
+  if shape is jtu.PYTHON_SCALAR_SHAPE:
+    return [t for t in dtypes if t in python_scalar_dtypes]
+  return dtypes
+
+def _shape_and_dtypes(shapes, dtypes):
+  for shape in shapes:
+    for dtype in _valid_dtypes_for_shape(shape, dtypes):
+      yield (shape, dtype)
+
 def _compatible_shapes(shape):
   if np.ndim(shape) == 0 or shape in scalar_shapes:
     return [shape]
@@ -457,6 +469,47 @@ class JaxNumpyReducerTests(jtu.JaxTestCase):
     with self.assertRaisesRegex(ValueError, r"duplicate value in 'axis': \\(0, 0\\)"):
       jnp.sum(jnp.arange(3), (0, 0))
 
+  @jtu.sample_product(
+    [dict(shape=shape, dtype=dtype, axis=axis, weights_shape=weights_shape)
+      for shape, dtype in _shape_and_dtypes(nonempty_shapes, number_dtypes)
+      for axis in list(range(-len(shape), len(shape))) + [None] + [tuple(range(len(shape)))]
+      # `weights_shape` is either `None`, same as the averaged axis, or same as
+      # that of the input
+      for weights_shape in ([None, shape] if axis is None or len(shape) == 1 or isinstance(axis, tuple)
+                            else [None, (shape[axis],), shape])
+    ],
+    keepdims=([False, True] if numpy_version >= (1, 23) else [None]),
+    returned=[False, True],
+  )
+  def testAverage(self, shape, dtype, axis, weights_shape, returned, keepdims):
+    rng = jtu.rand_default(self.rng())
+    kwds = dict(returned=returned)
+    if keepdims is not None:
+      kwds['keepdims'] = keepdims
+    if weights_shape is None:
+      np_fun = lambda x: np.average(x, axis, **kwds)
+      jnp_fun = lambda x: jnp.average(x, axis, **kwds)
+      args_maker = lambda: [rng(shape, dtype)]
+    else:
+      np_fun = lambda x, weights: np.average(x, axis, weights, **kwds)
+      jnp_fun = lambda x, weights: jnp.average(x, axis, weights, **kwds)
+      args_maker = lambda: [rng(shape, dtype), rng(weights_shape, dtype)]
+    np_fun = jtu.promote_like_jnp(np_fun, inexact=True)
+    tol = {dtypes.bfloat16: 2e-1, np.float16: 1e-2, np.float32: 1e-5,
+           np.float64: 1e-12, np.complex64: 1e-5}
+    check_dtypes = shape is not jtu.PYTHON_SCALAR_SHAPE and numpy_version >= (1, 22)
+    if numpy_version == (1, 23, 0) and keepdims and weights_shape is not None and axis is not None:
+      # Known failure: https://github.com/numpy/numpy/issues/21850
+      pass
+    else:
+      try:
+        self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
+                                check_dtypes=check_dtypes, tol=tol)
+      except ZeroDivisionError:
+        self.skipTest("don't support checking for ZeroDivisionError")
+    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=check_dtypes,
+                          rtol=tol, atol=tol)
+
   @parameterized.named_parameters(
       jtu.cases_from_list(
         {"testcase_name":
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 748d15a77..4bc37b1fd 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -44,7 +44,6 @@ from jax._src import device_array
 from jax._src import dtypes
 from jax._src import test_util as jtu
 from jax._src.lax import lax as lax_internal
-from jax._src.numpy.lax_numpy import _promote_dtypes, _promote_dtypes_inexact
 from jax._src.numpy.util import _parse_numpydoc, ParsedDoc, _wraps
 from jax._src.util import prod, safe_zip
 from jax._src import array
@@ -139,21 +138,6 @@ def _shapes_are_equal_length(shapes):
   return all(len(shape) == len(shapes[0]) for shape in shapes[1:])
 
 
-def _promote_like_jnp(fun, inexact=False):
-  """Decorator that promotes the arguments of `fun` to `jnp.result_type(*args)`.
-
-  jnp and np have different type promotion semantics; this decorator allows
-  tests make an np reference implementation act more like an jnp
-  implementation.
-  """
-  _promote = _promote_dtypes_inexact if inexact else _promote_dtypes
-  def wrapper(*args, **kw):
-    flat_args, tree = tree_util.tree_flatten(args)
-    args = tree_util.tree_unflatten(tree, _promote(*flat_args))
-    return fun(*args, **kw)
-  return wrapper
-
-
 class LaxBackedNumpyTests(jtu.JaxTestCase):
   """Tests for LAX-backed Numpy implementation."""
 
@@ -1447,7 +1431,7 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     wrapped_axis = axis % len(base_shape)
     shapes = [base_shape[:wrapped_axis] + (size,) + base_shape[wrapped_axis+1:]
               for size, _ in zip(itertools.cycle([3, 1, 4]), arg_dtypes)]
-    @_promote_like_jnp
+    @jtu.promote_like_jnp
     def np_fun(*args, dtype=dtype):
       dtype = dtype or args[0].dtype
       args = [x if x.dtype != jnp.bfloat16 else x.astype(np.float32)
@@ -1692,7 +1676,7 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
   def testRepeat(self, axis, shape, dtype, repeats, fixed_size):
     rng = jtu.rand_default(self.rng())
     np_fun = lambda arg: np.repeat(arg, repeats=repeats, axis=axis)
-    np_fun = _promote_like_jnp(np_fun)
+    np_fun = jtu.promote_like_jnp(np_fun)
     if fixed_size:
       total_repeat_length = np.repeat(np.zeros(shape), repeats, axis).shape[axis or 0]
       jnp_fun = lambda arg, rep: jnp.repeat(arg, repeats=rep, axis=axis,
@@ -2388,7 +2372,7 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
       args_maker = lambda: [np.array([rng(shape, dtype) for dtype in dtypes])]
     else:
       args_maker = lambda: [[rng(shape, dtype) for dtype in dtypes]]
-    np_fun = _promote_like_jnp(np.column_stack)
+    np_fun = jtu.promote_like_jnp(np.column_stack)
     jnp_fun = jnp.column_stack
     with jtu.strict_promotion_if_dtypes_match(dtypes):
       self._CheckAgainstNumpy(jnp_fun, np_fun, args_maker)
@@ -2417,9 +2401,9 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
       args_maker = lambda: [[rng(shape, dtype) for dtype in dtypes]]
 
     if numpy_version < (1, 24):
-      np_fun = _promote_like_jnp(lambda *args: np.stack(*args, axis=axis).astype(out_dtype))
+      np_fun = jtu.promote_like_jnp(lambda *args: np.stack(*args, axis=axis).astype(out_dtype))
     else:
-      np_fun = _promote_like_jnp(partial(np.stack, axis=axis, dtype=out_dtype, casting='unsafe'))
+      np_fun = jtu.promote_like_jnp(partial(np.stack, axis=axis, dtype=out_dtype, casting='unsafe'))
 
     jnp_fun = partial(jnp.stack, axis=axis, dtype=out_dtype)
     with jtu.strict_promotion_if_dtypes_match(dtypes):
@@ -2447,9 +2431,9 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
       args_maker = lambda: [[rng(shape, dtype) for dtype in dtypes]]
 
     if numpy_version < (1, 24) or op == "dstack":
-      np_fun = _promote_like_jnp(lambda *args: getattr(np, op)(*args).astype(out_dtype))
+      np_fun = jtu.promote_like_jnp(lambda *args: getattr(np, op)(*args).astype(out_dtype))
     else:
-      np_fun = partial(_promote_like_jnp(getattr(np, op)), dtype=out_dtype)
+      np_fun = partial(jtu.promote_like_jnp(getattr(np, op)), dtype=out_dtype)
 
     jnp_fun = partial(getattr(jnp, op), dtype=out_dtype)
     with jtu.strict_promotion_if_dtypes_match(dtypes):
@@ -2904,47 +2888,6 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
     self._CompileAndCheck(jnp_fun, args_maker)
 
-  @jtu.sample_product(
-    [dict(shape=shape, dtype=dtype, axis=axis, weights_shape=weights_shape)
-      for shape, dtype in _shape_and_dtypes(nonempty_shapes, number_dtypes)
-      for axis in list(range(-len(shape), len(shape))) + [None]
-      # `weights_shape` is either `None`, same as the averaged axis, or same as
-      # that of the input
-      for weights_shape in ([None, shape] if axis is None or len(shape) == 1
-                            else [None, (shape[axis],), shape])
-    ],
-    keepdims=([False, True] if numpy_version >= (1, 23) else [None]),
-    returned=[False, True],
-  )
-  def testAverage(self, shape, dtype, axis, weights_shape, returned, keepdims):
-    rng = jtu.rand_default(self.rng())
-    kwds = dict(returned=returned)
-    if keepdims is not None:
-      kwds['keepdims'] = keepdims
-    if weights_shape is None:
-      np_fun = lambda x: np.average(x, axis, **kwds)
-      jnp_fun = lambda x: jnp.average(x, axis, **kwds)
-      args_maker = lambda: [rng(shape, dtype)]
-    else:
-      np_fun = lambda x, weights: np.average(x, axis, weights, **kwds)
-      jnp_fun = lambda x, weights: jnp.average(x, axis, weights, **kwds)
-      args_maker = lambda: [rng(shape, dtype), rng(weights_shape, dtype)]
-    np_fun = _promote_like_jnp(np_fun, inexact=True)
-    tol = {dtypes.bfloat16: 2e-1, np.float16: 1e-2, np.float32: 1e-5,
-           np.float64: 1e-12, np.complex64: 1e-5}
-    check_dtypes = shape is not jtu.PYTHON_SCALAR_SHAPE and numpy_version >= (1, 22)
-    if numpy_version == (1, 23, 0) and keepdims and weights_shape is not None and axis is not None:
-      # Known failure: https://github.com/numpy/numpy/issues/21850
-      pass
-    else:
-      try:
-        self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker,
-                                check_dtypes=check_dtypes, tol=tol)
-      except ZeroDivisionError:
-        self.skipTest("don't support checking for ZeroDivisionError")
-    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=check_dtypes,
-                          rtol=tol, atol=tol)
-
   @jtu.sample_product(
     [dict(arg=arg, dtype=dtype, ndmin=ndmin)
       for arg, dtypes in [
@@ -4014,7 +3957,7 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     rng = jtu.rand_default(self.rng())
     args_maker = self._GetArgsMaker(rng, shapes, dtypes)
     def np_fun(cond, x, y):
-      return _promote_like_jnp(partial(np.where, cond))(x, y)
+      return jtu.promote_like_jnp(partial(np.where, cond))(x, y)
     with jtu.strict_promotion_if_dtypes_match(dtypes):
       self._CheckAgainstNumpy(np_fun, jnp.where, args_maker)
       self._CompileAndCheck(jnp.where, args_maker)
@@ -4978,7 +4921,7 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
       tol = {np.complex64: 1e-5, np.complex128: 1e-14}
 
     with jtu.strict_promotion_if_dtypes_match(dtypes):
-      self._CheckAgainstNumpy(_promote_like_jnp(np_op), jnp.logaddexp, args_maker, tol=tol)
+      self._CheckAgainstNumpy(jtu.promote_like_jnp(np_op), jnp.logaddexp, args_maker, tol=tol)
       self._CompileAndCheck(jnp.logaddexp, args_maker, rtol=tol, atol=tol)
 
   @jtu.sample_product(
@@ -5004,7 +4947,7 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
       tol = {np.complex64: 1e-5, np.complex128: 1e-14}
 
     with jtu.strict_promotion_if_dtypes_match(dtypes):
-      self._CheckAgainstNumpy(_promote_like_jnp(np_op), jnp.logaddexp2, args_maker, tol=tol)
+      self._CheckAgainstNumpy(jtu.promote_like_jnp(np_op), jnp.logaddexp2, args_maker, tol=tol)
       self._CompileAndCheck(jnp.logaddexp2, args_maker, rtol=tol, atol=tol)
 
   def testDefaultDtypes(self):