Repository: yt-project/unyt

Issue:
Should `apply_over_axes` "ignore" units?
This is a question more than a bug.

The current unyt implementation of numpy's `apply_over_axes` is:

```
@implements(np.apply_over_axes)
def apply_over_axes(func, a, axes):
    res = func(np.asarray(a), axes[0]) * a.units
    if len(axes) > 1:
        # this function is recursive by nature,
        # here we intentionally do not call the base _implementation
        return np.apply_over_axes(func, res, axes[1:])
    else:
        return res
```

This assumes that the function `func` preserves the units of the input `a`. It looks like it's necessary for `func` to return values with common units, but I guess those don't strictly need to be the same as those of `a`? As a trivial example, shouldn't:
```
def f(x):
    return x ** 2
```
be allowed? In this case the result will be incorrect. Is the following implementation enough?

```
@implements(np.apply_over_axes)
def apply_over_axes(func, a, axes):
    res = func(a, axes[0])
    if len(axes) > 1:
        # this function is recursive by nature,
        # here we intentionally do not call the base _implementation
        return np.apply_over_axes(func, res, axes[1:])
    else:
        return res
```

Discussion:
kyleaoman: pre-commit.ci autofix
kyleaoman: Did not know about that syntax for parametrize - cool!

Patch:
diff --git a/unyt/_array_functions.py b/unyt/_array_functions.py
index c8e92f9..03a4875 100644
--- a/unyt/_array_functions.py
+++ b/unyt/_array_functions.py
@@ -850,7 +850,7 @@ def savetxt(fname, X, *args, **kwargs):
 
 @implements(np.apply_over_axes)
 def apply_over_axes(func, a, axes):
-    res = func(np.asarray(a), axes[0]) * a.units
+    res = func(a, axes[0])
     if len(axes) > 1:
         # this function is recursive by nature,
         # here we intentionally do not call the base _implementation


Test patch:
diff --git a/unyt/tests/test_array_functions.py b/unyt/tests/test_array_functions.py
index fbabac1..a8b582c 100644
--- a/unyt/tests/test_array_functions.py
+++ b/unyt/tests/test_array_functions.py
@@ -1654,11 +1654,16 @@ def test_apply_along_axis():
     assert ret.units == cm**2
 
 
-def test_apply_over_axes():
+@pytest.mark.parametrize("axes, expected_units", [((0, 1), cm**4), ((0,), cm**2)])
+def test_apply_over_axes(axes, expected_units):
+    # the user-supplied function must be trusted to treat units
+    # sensibly (mainly that it doesn't give a mix of units across
+    # the resulting array), but we can check that units are
+    # propagated correctly for well-behaved functions.
     a = np.eye(3) * cm
-    ret = np.apply_over_axes(lambda x, axis: x * cm, a, (0, 1))
-    assert type(ret) is unyt_array
-    assert ret.units == cm**3
+    ret = np.apply_over_axes(lambda x, axis: x[axis] ** 2, a, axes)
+    assert isinstance(ret, unyt_array)  # could be unyt_quantity
+    assert ret.units == expected_units
 
 
 def test_array_equal():