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
```

Hints from issue comments:
@neutrinoceros: > shouldn't:
>```python
> def f(x):
>    return x ** 2
>```
>be allowed?

It should actually not: `np.apply_over_axis` only accepts function that have 2 arguments, where the second one must be an axis. 

> In this case the result will be incorrect.

I'm not sure what you meant here.@kyleaoman: Hm, let me try to be more explicit:
```
>>> def f(x, axis):
...     return x[axis] ** 2
... 
>>> np.apply_over_axes(f, np.arange(9).reshape((3, 3)) * u.cm, (0, ))
unyt_array([0, 1, 4], 'cm')
```
Shouldn't the result have `cm**2`?@neutrinoceros: indeed it should. Thanks for clarifying. I now agree that unyt's implementation of `apply_over_axis` shouldn't be concerned with units, as it's both out of its hands and not its responsibility. I'll accept a PR with your proposed change, which makes a lot more sense. Thanks !

Solution patch:
diff --git a/unyt/_array_functions.py b/unyt/_array_functions.py
index c8e92f97..03a4875d 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 fbabac1d..a8b582c6 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():