Merge branch 'main' into maxymnaumchyk/add-to-from-tensorflow-functions

.all-contributorsrc

@@ -501,6 +501,15 @@
       "contributions": [
         "maintenance"
       ]
+    },
+    {
+      "login": "nj-vs-vh",
+      "name": "Igor Vaiman",
+      "avatar_url": "https://avatars.githubusercontent.com/u/30616208?v=4",
+      "profile": "https://nj-vs-vh.name/",
+      "contributions": [
+        "code"
+      ]
     }
   ],
   "contributorsPerLine": 7,

README.md

@@ -232,6 +232,7 @@ Thanks especially to the gracious help of Awkward Array contributors (including
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/maxymnaumchyk"><img src="https://avatars.githubusercontent.com/u/70752300?v=4?s=100" width="100px;" alt="maxymnaumchyk"/><br /><sub><b>maxymnaumchyk</b></sub></a><br /><a href="https://github.com/scikit-hep/awkward/commits?author=maxymnaumchyk" title="Code">💻</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://tacaswell.github.io"><img src="https://avatars.githubusercontent.com/u/199813?v=4?s=100" width="100px;" alt="Thomas A Caswell"/><br /><sub><b>Thomas A Caswell</b></sub></a><br /><a href="#maintenance-tacaswell" title="Maintenance">🚧</a></td>
       <td align="center" valign="top" width="14.28%"><a href="http://www.nijho.lt"><img src="https://avatars.githubusercontent.com/u/6897215?v=4?s=100" width="100px;" alt="Bas Nijholt"/><br /><sub><b>Bas Nijholt</b></sub></a><br /><a href="#maintenance-basnijholt" title="Maintenance">🚧</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://nj-vs-vh.name/"><img src="https://avatars.githubusercontent.com/u/30616208?v=4?s=100" width="100px;" alt="Igor Vaiman"/><br /><sub><b>Igor Vaiman</b></sub></a><br /><a href="https://github.com/scikit-hep/awkward/commits?author=nj-vs-vh" title="Code">💻</a></td>
     </tr>
   </tbody>
 </table>

src/awkward/operations/ak_mean.py

@@ -225,8 +225,8 @@ def _impl(x, weight, axis, keepdims, mask_identity, highlevel, behavior, attrs):
             sumw = ak.operations.ak_sum._impl(
                 x * 0 + weight,
                 axis,
-                keepdims,
-                mask_identity,
+                keepdims=True,
+                mask_identity=True,
                 highlevel=True,
                 behavior=ctx.behavior,
                 attrs=ctx.attrs,

tests/test_3285_ak_mean_weighted_row_wise.py

@@ -0,0 +1,65 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+
+from __future__ import annotations
+
+import math
+
+import pytest
+
+import awkward as ak
+
+
+@pytest.mark.parametrize(
+    "keepdims, expected_result",
+    [
+        pytest.param(False, ak.Array([2.25, 6.5])),
+        pytest.param(True, ak.Array([[2.25], [6.5]])),
+    ],
+)
+def test_keepdims(keepdims: bool, expected_result: ak.Array):
+    data = ak.Array(
+        [
+            [1, 2, 3],
+            [4, 7],
+        ]
+    )
+    weight = ak.Array(
+        [
+            [1, 1, 2],
+            [1, 5],
+        ]
+    )
+    assert ak.all(
+        ak.mean(data, weight=weight, axis=1, keepdims=keepdims) == expected_result
+    )
+
+
+@pytest.mark.parametrize(
+    "mask_identity, expected_result",
+    [
+        pytest.param(False, ak.Array([1.5, math.nan, 8])),
+        pytest.param(True, ak.Array([1.5, None, 8])),
+    ],
+)
+def test_mask_identity(mask_identity: bool, expected_result: ak.Array):
+    data = ak.Array(
+        [
+            [1, 2],
+            [],
+            [6, 9],
+        ]
+    )
+    weight = ak.Array(
+        [
+            [1, 1],
+            [],
+            [1, 2],
+        ]
+    )
+    result = ak.mean(data, weight=weight, axis=1, mask_identity=mask_identity)
+    assert result[0] == expected_result[0]
+    if mask_identity:
+        assert result[1] is None
+    else:
+        assert math.isnan(result[1])  # NaN is not equal to itself per IEEE!
+    assert result[2] == expected_result[2]