Repository: scikit-hep/awkward

Issue:
ak.mean with `keepdims=True` is returning arrays of the wrong length
### Discussed in https://github.com/scikit-hep/awkward/discussions/2726

<div type='discussions-op-text'>

<sup>Originally posted by **WWWonderer** September 27, 2023</sup>
For the code below:

```
arr = ak.Array(
    [
        [
            [1, 2], 
            [1]
        ],
        [
            [1], 
            [1], 
            [1]
        ]
    ]
)

print(f'mean: {ak.mean(arr)}') # prints mean: 1.1666666666666667
print(f'std: {ak.std(arr)}') # error
```
The mean works fine while the std gives me error with following stack trace:
```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
[/usr/local/lib/python3.10/dist-packages/awkward/_dispatch.py](https://localhost:8080/#) in dispatch(*args, **kwargs)
     59                 try:
---> 60                     next(gen_or_result)
     61                 except StopIteration as err:

19 frames
ValueError: cannot broadcast nested list (in compiled code: https://github.com/scikit-hep/awkward/blob/awkward-cpp-23/awkward-cpp/src/cpu-kernels/awkward_ListArray_broadcast_tooffsets.cpp#L27)

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
[/usr/local/lib/python3.10/dist-packages/awkward/_errors.py](https://localhost:8080/#) in handle_exception(self, cls, exception)
     80             self.decorate_exception(cls, exception)
     81         else:
---> 82             raise self.decorate_exception(cls, exception)
     83 
     84     def decorate_exception(self, cls: type[E], exception: E) -> E:

ValueError: cannot broadcast nested list (in compiled code: https://github.com/scikit-hep/awkward/blob/awkward-cpp-23/awkward-cpp/src/cpu-kernels/awkward_ListArray_broadcast_tooffsets.cpp#L27)

This error occurred while calling

    ak.std(
        <Array [[[1, 2], [1]], [[1], ...]] type='2 * var * var * int64'>
    )
```
This seems very counterintuitive to me as the std is nothing but some operations on mean. What can be done in mean in terms of broadcasting should also be doable in std. Is there a reason this is happening and how to solve this issue? Thanks!

Edit: 

I went a bit deeper into the codebase and found this in variance implementation which is used by std (https://github.com/scikit-hep/awkward/blob/main/src/awkward/operations/ak_var.py):
```
            sumwxx = ak.operations.ak_sum._impl(
                (x - xmean) ** 2, # maybe here is the broadcasting??? 
                axis,
                keepdims=True,
                mask_identity=True,
                highlevel=True,
                behavior=None,
            )
```
In the implementation of mean (https://github.com/scikit-hep/awkward/blob/main/src/awkward/operations/ak_mean.py#L156) only x is passed into ak_sum:
```
            sumwx = ak.operations.ak_sum._impl(
                x,
                axis,
                keepdims=True,
                mask_identity=True,
                highlevel=True,
                behavior=None,
            )
```
Maybe this is the cause of the issue? But is there a way to fix it? 
</div>

--------------

I looked into the implementation of `ak.std`, and found the same issue in `ak.var`. The step where it fails is when we subtract the mean from each value:

```python
xmean = ak.mean(arr, keepdims=True)
arr - xmean   # fails with "can't broadcast"
```

where

```python
>>> arr
<Array [[[1, 2], [1]], [[1], [1], [1]]] type='2 * var * var * int64'>
>>> xmean
<Array [[[1.17]]] type='1 * var * 1 * float64'>
```

The problem is not with the `var * int64` versus `1 * float64`, it's with the `2 *` versus `1 *`. It doesn't look like a problem with the definition of broadcasting, but with the output of `ak.mean` with `keepdims=True`. It should be length-2, not length-1.

(I don't know how this didn't get caught in tests! My mistake, of course, since I wrote the `ak.mean`, `ak.var`, etc. suite.)


Discussion:
The cause of this is that `keepdims=True` for `axis=None` doesn't introduce regular nodes, unlike `axis!=None` which replaces the node at `axis` with a regular one. 

This makes sense; unlike `axis=None`, `axis=N` only operates upon a single dimension. We _probably_ want to replace list nodes with regular nodes in `_remove_structure`.

Resolving patch:
diff --git a/src/awkward/_do.py b/src/awkward/_do.py
--- a/src/awkward/_do.py
+++ b/src/awkward/_do.py
@@ -201,6 +201,7 @@ def remove_structure(
     drop_nones: bool = True,
     keepdims: bool = False,
     allow_records: bool = False,
+    list_to_regular: bool = False,
 ):
     if isinstance(layout, Record):
         return remove_structure(
@@ -224,6 +225,7 @@ def remove_structure(
                 "drop_nones": drop_nones,
                 "keepdims": keepdims,
                 "allow_records": allow_records,
+                "list_to_regular": list_to_regular,
             },
         )
         return tuple(arrays)
@@ -272,6 +274,7 @@ def reduce(
             drop_nones=False,
             keepdims=keepdims,
             allow_records=True,
+            list_to_regular=True,
         )
 
         if len(parts) > 1:
diff --git a/src/awkward/contents/listoffsetarray.py b/src/awkward/contents/listoffsetarray.py
--- a/src/awkward/contents/listoffsetarray.py
+++ b/src/awkward/contents/listoffsetarray.py
@@ -2065,18 +2065,34 @@ def _remove_structure(self, backend, options):
             content = self._content[self._offsets[0] : self._offsets[-1]]
             contents = content._remove_structure(backend, options)
             if options["keepdims"]:
-                return [
-                    ListOffsetArray(
-                        Index64(
-                            backend.index_nplike.asarray(
-                                [0, backend.index_nplike.shape_item_as_index(c.length)]
-                            )
-                        ),
-                        c,
-                        parameters=self._parameters,
-                    )
-                    for c in contents
-                ]
+                if options["list_to_regular"]:
+                    return [
+                        ak.contents.RegularArray(
+                            c,
+                            size=c.length,
+                            zeros_length=1,
+                            parameters=self._parameters,
+                        )
+                        for c in contents
+                    ]
+                else:
+                    return [
+                        ListOffsetArray(
+                            Index64(
+                                backend.index_nplike.asarray(
+                                    [
+                                        0,
+                                        backend.index_nplike.shape_item_as_index(
+                                            c.length
+                                        ),
+                                    ]
+                                )
+                            ),
+                            c,
+                            parameters=self._parameters,
+                        )
+                        for c in contents
+                    ]
             else:
                 return contents
 
diff --git a/src/awkward/contents/regulararray.py b/src/awkward/contents/regulararray.py
--- a/src/awkward/contents/regulararray.py
+++ b/src/awkward/contents/regulararray.py
@@ -1085,7 +1085,7 @@ def _reduce_next(
                 )
             )
 
-            if self._size > 0:
+            if self._size is not unknown_length and self._size > 0:
                 nextstarts = ak.index.Index64(
                     index_nplike.arange(0, nextlen, self._size),
                     nplike=index_nplike,
@@ -1151,7 +1151,11 @@ def _reduce_next(
                     )
 
                     trimmed = outcontent.content._getitem_range(start, stop)
-                    assert len(trimmed) == self._size * outcontent.length
+                    assert (
+                        trimmed.length is unknown_length
+                        or outcontent.length is unknown_length
+                        or trimmed.length == self._size * outcontent.length
+                    )
 
                     outcontent = ak.contents.RegularArray(
                         trimmed,


Test patch:
diff --git a/tests/test_2727_remove_structure_regular.py b/tests/test_2727_remove_structure_regular.py
new file mode 100644
--- /dev/null
+++ b/tests/test_2727_remove_structure_regular.py
@@ -0,0 +1,29 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import numpy as np
+import pytest
+
+import awkward as ak
+
+
+def test():
+    array = ak.Array([[[1, 2], [1]], [[1], [1], [1]]])
+    assert ak.almost_equal(
+        ak.mean(array, axis=None, keepdims=True),
+        ak.contents.RegularArray(
+            ak.contents.RegularArray(
+                ak.contents.NumpyArray(
+                    np.mean([1, 2, 1, 1, 1, 1], dtype=np.float64, keepdims=True)
+                ),
+                size=1,
+            ),
+            size=1,
+        ),
+        check_regular=True,
+    )
+
+
+def test_original_problem():
+    array = ak.Array([[[1, 2], [1]], [[1], [1], [1]]])
+    assert ak.mean(array) == pytest.approx(1.1666666666666667)
+    assert ak.var(array) == pytest.approx(0.13888888888888887)