Repository: yt-project/unyt

Issue:
Incorrect units when using `np.prod` with an `axis` kwarg
* unyt version: 3.0.3.post34+gecd71ea
* Python version: 3.12
* Operating System: Linux (Fedora 40)

### Description

There is a bug in the unyt implementation of the numpy `prod` function. It does not correctly set the units in the returned array when an axis keyword is passed.

### What I Did

```
>>> import unyt as u
>>> np.prod(np.arange(1,10).reshape((3,3)) * u.cm,axis=0)
unyt_array([ 28,  80, 162], 'cm**9')
```
Expected units are `cm**3` in this case.

Hints from issue comments:

Solution patch:
diff --git a/unyt/_array_functions.py b/unyt/_array_functions.py
index fa2c478c..790209af 100644
--- a/unyt/_array_functions.py
+++ b/unyt/_array_functions.py
@@ -535,7 +535,8 @@ def copyto(dst, src, *args, **kwargs):
 
 @implements(np.prod)
 def prod(a, *args, **kwargs):
-    return np.prod._implementation(np.asarray(a), *args, **kwargs) * a.units**a.size
+    res = np.prod._implementation(np.asarray(a), *args, **kwargs)
+    return res * a.units ** (a.size // res.size)
 
 
 @implements(np.var)


Test patch:
diff --git a/unyt/tests/test_array_functions.py b/unyt/tests/test_array_functions.py
index bfd900ab..4ae97a04 100644
--- a/unyt/tests/test_array_functions.py
+++ b/unyt/tests/test_array_functions.py
@@ -1291,6 +1291,18 @@ def test_scalar_reducer(func, expected_units):
     assert y.units == expected_units
 
 
+@pytest.mark.parametrize("func", [np.prod, np.nanprod])
+def test_prod_with_axis(func):
+    x = [
+        [1, 2, 3],
+        [4, 5, 6],
+        [7, 8, 9],
+    ] * cm
+    y = func(x, axis=0)
+    assert type(y) is unyt_array
+    assert y.units == cm**3
+
+
 @pytest.mark.parametrize(
     "func",
     [