Repository: hgrecco/pint

Issue:
Quantity.prod() result unit is incorrect with axis or where argument
Right now, the `.prod()` method assumes the full size of the input is collapsed in the result. This gives incorrect results when the `axis` or `where` arguments are supplied, as seen below:

```python
import pint

ureg = pint.UnitRegistry()

q = [[1, 2], [3, 4]] * ureg.m

print(q.prod())
print(q.prod(axis=0))
print(q.prod(where=[True, False]))
```
```
24 meter ** 4
[3 8] meter ** 4
3 meter ** 4
```

The unit on the first result where the full array is collapsed to a scalar is correct, but the other two results should have `meter ** 2` as the unit.

I don't have a good fix for this yet, since this is the same problem mentioned in https://github.com/hgrecco/pint/pull/764#issuecomment-523749008. If anyone has any suggestions for a performant way to determine how many unit multiplications occur given both `axis` and `where` arguments, please do let me know! Otherwise, I'll try to figure something out and get a PR in for this or include it alongside a `prod` implementation for `__array_function__`.

Discussion:

Patch:
diff --git a/pint/numpy_func.py b/pint/numpy_func.py
index 8501a7e..7188071 100644
--- a/pint/numpy_func.py
+++ b/pint/numpy_func.py
@@ -670,6 +670,28 @@ def _all(a, *args, **kwargs):
         raise ValueError("Boolean value of Quantity with offset unit is ambiguous.")
 
 
+@implements("prod", "function")
+def _prod(a, *args, **kwargs):
+    arg_names = ("axis", "dtype", "out", "keepdims", "initial", "where")
+    all_kwargs = dict(**dict(zip(arg_names, args)), **kwargs)
+    axis = all_kwargs.get("axis", None)
+    where = all_kwargs.get("where", None)
+
+    if axis is not None and where is not None:
+        raise ValueError("passing axis and where is not supported")
+
+    result = np.prod(a._magnitude, *args, **kwargs)
+
+    if axis is not None:
+        exponent = a.size // result.size
+        units = a.units ** exponent
+    elif where is not None:
+        exponent = np.asarray(where, dtype=np.bool_).sum()
+        units = a.units ** exponent
+
+    return units._REGISTRY.Quantity(result, units)
+
+
 # Implement simple matching-unit or stripped-unit functions based on signature
 
 


Test patch:
diff --git a/pint/testsuite/test_numpy.py b/pint/testsuite/test_numpy.py
index fd12bf5..594bde0 100644
--- a/pint/testsuite/test_numpy.py
+++ b/pint/testsuite/test_numpy.py
@@ -286,6 +286,16 @@ class TestNumpyMathematicalFunctions(TestNumpyMethods):
     def test_prod(self):
         self.assertEqual(self.q.prod(), 24 * self.ureg.m ** 4)
 
+    @helpers.requires_array_function_protocol()
+    def test_prod_numpy_func(self):
+        axis = 0
+        where = [[True, False], [True, True]]
+
+        self.assertQuantityEqual(np.prod(self.q, axis=axis), [3, 8] * self.ureg.m ** 2)
+        self.assertQuantityEqual(np.prod(self.q, where=where), 12 * self.ureg.m ** 3)
+
+        self.assertRaises(ValueError, np.prod, self.q, axis=axis, where=where)
+
     def test_sum(self):
         self.assertEqual(self.q.sum(), 10 * self.ureg.m)
         self.assertQuantityEqual(self.q.sum(0), [4, 6] * self.ureg.m)