ak.sum produces incorrect structure for outer dimension

[agoose77]

Version of Awkward Array

56312ea

This is seen in both v1 and v2 in master

Description and code to reproduce

Given this layout:

x = ak._v2.Array([
    [
        [
            [1, 2, 3],
            [4, 0, 0]
        ],
        [
            [5, 6, 7],
            [8, 0, 0]
        ]
    ],
    [
        [
            [9, 10, 11],
            [12, 0, 0]
        ],
        [
            [13, 14, 15],
            [16, 0, 0]
        ]
        
    ]
])

Calling ak.sum(x, axis=0) does not match (in structure) the NumPy result:

>>> ak._v2.to_numpy(x).sum(axis=0).tolist()
[[[10, 12, 14], [16, 0, 0]], [[18, 20, 22], [24, 0, 0]]]
>>> ak._v2.sum(x, axis=0).tolist()
[[[10, 12, 14], []], [[16, 0, 0, 18, 20, 22], []]]

[agoose77]

I'm not yet familiar enough with how we use _reduce_next to understand why we run the kernel routines, but I am able to see that awkward_ListOffsetArray_reduce_nonlocal_outstartsstops_64 is misbehaving. The outstarts and outstops arguments should have the same length as outlength, but if you use a Python implementation of this kernel, it raises an IndexError when an out-of-bounds index is accessed:

outstarts = np.zeros(4, dtype=np.int64)
outstops = np.zeros(4, dtype=np.int64)

distincts, distincts_length, gaps, outlength = (
    np.array([0, 0, 0, 2, 2, 2, 1, 1, 1, 3, 3, 3], dtype=np.int64),
    12,
    np.array([1, 2, 1, 140187859207408], dtype=np.int64),
    4,
)

def awkward_ListOffsetArray_reduce_nonlocal_outstartsstops_64(
    outstarts, outstops, distincts, lendistincts, gaps, outlength
):
    maxcount = lendistincts if outlength == 0 else lendistincts // outlength

    j = 0
    k = 0
    maxdistinct = -1
    lasti = -1
    for i in range(lendistincts):
        j_0 = j
        if maxdistinct < distincts[i]:
            maxdistinct = distincts[i]

            extra = (i - lasti) // maxcount
            lasti = i
            numgappy = gaps[j]
            if numgappy < extra:
                numgappy = extra

            # Create gaps for blank sublists with start=stop
            for gappy in range(numgappy):
                outstarts[k] = i
                outstops[k] = i
                k = k + 1

            j = j + 1

        if distincts[i] != -1:
            outstops[k - 1] = i + 1

    # Fill the remaining start-stops to produce empty lists
    for k in range(k, outlength):
        outstarts[k] = lendistincts + 1
        outstops[k] = lendistincts + 1
        


awkward_ListOffsetArray_reduce_nonlocal_outstartsstops_64(
    outstarts, outstops, distincts, distincts_length, gaps, outlength
)

[agoose77]

I don't think this is actually the cause of this problem though (because we never read these o-o-b values). Instead, something looks funky with the distincts ([0, 0, 0, 2, 2, 2, 1, 1, 1, 3, 3, 3]) and the algorithm that uses them here.

[jpivarski]

Calling ak.sum(x, axis=0) does not match (in structure) the NumPy result:

>>> ak._v2.to_numpy(x).sum(axis=0).tolist()
[[[10, 12, 14], [16, 0, 0]], [[18, 20, 22], [24, 0, 0]]]
>>> ak._v2.sum(x, axis=0).tolist()
[[[10, 12, 14], []], [[16, 0, 0, 18, 20, 22], []]]

This does look wrong. Does it work in v1? Cases like this were extensively tested (and I thought I had ported all of those tests to v2). Since it's a structure thing, it shouldn't matter which reducer we use, and I find that prod with prime numbers is a good diagnostic, since there's only one way to obtain a given product. (WIth sum, you'd have to use powers of 2.)

If it is a v1/v2 difference, the issue wouldn't be in the kernel itself, since that's the same for both. It would be a matter of how it's called and what Indexes it's writing into.

[agoose77]

@jpivarski it doesn't work in v1 either. As I understand our reducers, and based upon what you've said, the _reduce_next ultimately assembles the elements in such a way that the reducer only has to walk over a array of monotonically increasing parents. It doesn't look like this is true in this case: in a simpler example, the result of awkward_ListOffsetArray_reduce_nonlocal_preparenext_64 is

 nextparents=[0 0 2 2 1 1 3 3], nextcarry=[0 2 4 6 1 3 5 7]

which means that later kernel functions that appear to assume parents are monotonic fails (which is why I suspected awkward_ListOffsetArray_reduce_nonlocal_outstartsstops_64 at first).

[agoose77]

This simpler example is

import awkward as ak
x = ak._v2.Array([
    [
        [
            [1],
            [4]
        ],
        [
            [5],
            [8]
        ]
    ],
    [
        [
            [9],
            [12]
        ],
        [
            [13],
            [16]
        ]
        
    ]
])

print(ak._v2.sum(x, axis=1))
print(ak._v2.to_numpy(x).sum(axis=1))

[agoose77]

I modified listoffsetarray.py https://github.com/scikit-hep/awkward-1.0/blob/da762fa814531806d7d95a4ac93b43bccefe9dfe/src/awkward/_v2/contents/listoffsetarray.py#L867 locally to order carry such that nextparents was monotonic:

if depth == 3:
    nextparents[:] = [0, 0, 1, 1, 2, 2, 3, 3]
    nextcarry[:] = [0, 2, 1, 3, 4, 6, 5, 7]

and this then produces the correct result, making me think it's the aforementioned kernel computing nextparents and nextcarry incorrectly.

[jpivarski]

nextparents=[0 0 2 2 1 1 3 3]

It's more likely that I'm misremembering the monotonicity of parents than that the code that generated the above is wrong. It's possible that I'm misremembering the monotonicity of parents. Maybe the above is right.

Let me backpedal on that statement: at least integer values are guaranteed to be contiguous—i.e. the second 2 can't appear at the end of that array. The whole thing was designed around making sure that parallel jagged reduction would work, and it wouldn't work unless values in the parents are contiguous. In fact, this cpu-kernel is "sequential algorithm 2: scatter from parents" in the linked notebook, which also relies on the parents being value-contiguous.

Argmin and argmax need a starts to go along with this, and it's possible to deal with non-contiguous parents with starts that are out of order. In the above case, the starts could be

[0 4 2 6]

Edit: for the moment, I'm going to be non-committal about whether parents and starts need to be monotonic, in addition to parents needing to be value-contiguous.

[jpivarski]

The _reduce_next method, here in v1 and here in v2, has two main cases: the recursion depth is above the level chosen by axis, and the recursion depth is below the level chosen by axis. The case axis=-1 is "below" all the way down to the NumpyArray. The "above" case is selected by if not branch and negaxis == depth here in v1 and here in v2. The "below" case is selected by else here in v1 and here in v2, to help with navigation. It looks like I used the words "nonlocal" for "above" and "local" for "below" in the kernel names.

The "below" case is the easier one, and for axis=-1, the parents will definitely be monatonic. Here, the problem is just one of reducing together (e.g. adding up) array elements that are next to each other, though there's an arbitrary number of them in each group. np.ufunc.reduceat does this using only starts, and that's how it was implemented in Awkward 0. In Awkward 1, it was implemented with parents because (a) the "sequential algorithm 2: scatter from parents" was found to be faster, (b) it's the only way to do it on a vectorized/parallel device like a GPU, and (c) because this mixes better with the "above" case.

The "above" case was motivated by scikit-hep/awkward-0.x#166. The implementation strategy started by just trying to reproduce NumPy for fixed list lengths by reordering the next level down (reordering Index is nextcarry; reordered data is nextcontent) and keeping track of where each element came from with the next parents. Here is where I thought the nextparents would be monotonic, because the nextcontent is rearranged in such a way as to make it so. As I said above, I'm not entirely sure about that.

Things like "maxcount," "findgaps," and "offsetscopy" are complications for dealing with variable list lengths, beyond NumPy. Since you saw an error in a purely NumPy-like case, those aren't relevant to the problem. Similarly, the kernels that manage the starts are not relevant because starts are only needed for argmin and argmax.

The testing started in tests/test_0115-generic-reducer-operation.py. Notice that they start by dealing with the NumPy-like case before moving on to issues with "gaps" and such. The NumPy-like tests all work on a 3D array

>>> np.array(primes)[:2*3*5].reshape(2, 3, 5)
array([[[  2,   3,   5,   7,  11],
        [ 13,  17,  19,  23,  29],
        [ 31,  37,  41,  43,  47]],

       [[ 53,  59,  61,  67,  71],
        [ 73,  79,  83,  89,  97],
        [101, 103, 107, 109, 113]]])

and test all possible axis depths: 0, 1, 2. Your minimal reproducer is a 4D array, and I can't get a reproducer with only 3 dimensions. I would have thought that tests that include "all below," "all above," and "one below, one above" would have exhausted the possibilities. However, your reproducer fails for axis=0 and axis=1, which is to say, all cases with "at least two below."

>>> ak.sum(ak.Array([[[[1], [4]], [[5], [8]]], [[[9], [12]], [[13], [16]]]]), axis=0).tolist()
[[[10], []], [[16, 18], []]]
>>> ak.sum(ak.Array([[[[1], [4]], [[5], [8]]], [[[9], [12]], [[13], [16]]]]), axis=1).tolist()
[[[6], []], [[12, 22], []]]

The axis=2 and axis=3 cases are fine.

Reduction does not treat RegularArrays any differently from List/ListOffsetArrays, except inasmuch as they preserve regularity in the output (a later correction, #447).

For what it's worth, the reducer algorithms were first developed outside the context of Awkward Arrays, in studies/reducers.py, to be able to make tests rapidly and avoid the formality of defining kernels right away. This file might be useful to see how the kernels work together without the special cases that have accumulated over the years. See, for instance,

https://github.com/scikit-hep/awkward-1.0/blob/8e682003331d0bbe493f4a8b3446f8cc081878b8/studies/reducers.py#L964-L1050

for a straight-through implementation of ListOffsetArray._reduce_next without kernel indirection.

One last thing: the only kernel that updates a specially prepared Index is awkward_ListOffsetArray_reduce_nonlocal_preparenext, which uses and changes-in-place offsetscopy. (That's why it's a copy.) The studies file shows why we want that.

[agoose77]

I've looked into this some more, and I now understand that this bug is a manifestation of parents having non well-defined semantics. Some kernels (findgaps and nonlocal_outstartstops) seem to assume that the parent array will only every increase or remain constant.

It looks like these kernels are computing gaps and distincts are incorrectly. If I patch the code here to set gaps[:] = [1, 1, 1, 1], and modify nonlocal_outstartstops to detect list boundaries with lastdistinct != distinct[i] then the result is correct.

I'd like to just clarify my understanding of what these arrays mean:

• distincts: compute densely mapped parents, e.g. [0, 1, 3, 4] -> [0, 1, -1, 2, 3]
• gaps: compute missing sublists, where 1 is a list boundary, 2 is a missing sublist, etc, e.g. [0, 1, 3, 4] -> [1, 1, 2, 1]

As far as the actual reduction is concerned, the order of parents doesn't matter - it just reduces and writes to the output. It's only when rebuilding the structure that we care about this order.

More concretely, consider this example:

x = ak._v2.Array([
    [
        [
            [1],
            [4]
        ],
        [
            [5],
            [8]
        ],
        [
            [6],
            [7]
        ]
    ],
    [
        [
            [9],
            [12]
        ],
        [
            [13],
            [16]
        ]

    ]
])

Assume this layout is only ListOffsetArrays and a final NumpyArray, then we have at the deepest level:

data = [1, 4, 5, 8, 6, 7, 9, 12, 13, 16]

At the nonlocal level, we want to produce a carry and nextparents such that the reducer can act in a single pass:

carry = [0, 2, 4, 1, 3, 5, 6, 8, 7, 9]
nextparents = [0, 0, 0, 1, 1, 1, 2, 2, 3, 3]

Whereas right now we have

carry = [0, 2, 4, 6, 8, 1, 3, 5, 7, 9]
nextparents = [0, 0, 0, 2, 2, 1, 1, 1, 3, 3]

If this is the case, I think we need parents to be ordered. We currently use that information to identify empty sub-lists, and we can't do that if this assumption does not hold. Requiring this ordering would also simplify the kernel used to compute the output start/stops, because we could use ListOffsetArray instead of ListArray

[jpivarski]

Especially given the fact that I don't remember now, it's entirely possible that I had different ideas about the monotonicity of parents at different points in the development process. I'm still surprised that 3D tests weren't enough to catch that, but maybe this only comes about with dual application of "above" plus at least one "below"...

A way to test this idea, and possibly as a long-term fix, is to put a nextparents, nextcarry reordering step right after those two Indexes are first created and before nextparents is assumed to be monotonic. (It would be better to do this at the level of reordering nextcarry than by carrying twice.) They don't have to be reordered in place, and probably shouldn't be.

The studies script should still run as-is. If you put one of your reproducing examples into it, you should see the same error there because it's a flaw in the logic, and that script never tested a 4D regular array. Adding the reordering step would be a lot faster and hackable in that script than it would be in the real thing—you'd get a confirmation that it works before spending the time needed to create a new kernel. What do you think?

[agoose77]

A way to test this idea, and possibly as a long-term fix, is to put a nextparents, nextcarry reordering step right after those two Indexes are first created and before nextparents is assumed to be monotonic.

This would work as a band-aid, I tested it earlier :)

j = self._nplike.argsort(nextparents)
nextcarry = ak._v2.index.Index64(self._nplike.asarray(nextcarry)[j])
nextparents = ak._v2.index.Index64(self._nplike.asarray(nextparents)[j])

I think we probably can avoid this by re-writing prepare_next so that it outputs nextcarry in ascending nextparents order. If parents is always in order, then it is possible to iterate such that we produce nextparents in order. I need to spend some time to implement this, but I'm not able to do that just at the moment!

[jpivarski]

There are so many index-rewrites in this code-path that adding one more as a band-aid wouldn't disturb the performance much. If the question is whether that would obfuscate the code more, that's a more serious concern, but I think a kernel that only reorders could be easier to read. It has the advantage that we can guarantee that it is not wrong only by reading it—in isolation, unconnected to anything else. Certainly none of the other kernels are going to break by having monotonic parents instead of non-monotonic parents, and C code that only permutes two array buffers should be easy to verify for correctness by eye.

If the band-aid can be applied in short order (within your time constraints), then that would be more valuable than doing a more thorough job later, with the risk of it not happening at all. Things do fall through the cracks, and this is a bug-fix. Bug-fixes are important! :)

If you can write down, here perhaps, what that kernel should look like, another one of us can implement the kernel (i.e. adding to YAML and cpu-kernels). @ianna and I have experience in that workflow. Oh! I see that the code you wrote is a complete v2 fix. We might even have an equivalent of np.argsort as a kernel for v1 already. (@ianna, are any of the kernels you created usable here?) At what point in the code do you insert it? Here in v1 and here in v2?

[agoose77]

At what point in the code do you insert it? Here in v1 and here in v2?

Yes to both!

Yeah, if we put a todo to just fix the prepare_next kernel then I think this would fix the bug for now :)

I understand prepare_next well enough to know that it's wrong, but I need to spend more time on it in order to reason about what it should look like in order to return nicer parents. I'm sure it's not too tricky, but just a bit out of my time right now!

[ianna]

Hmm... I managed to crash v1 with the following test:

    array = np.array([
        [
            [
                [1],
                [4]
            ],
            [
                [5],
                [8]
            ]
        ],
        [
            [
                [9],
                [12]
            ],
            [
                [13],
                [16]
            ]

        ]
    ])
    nplike = ak.nplike.of(array)

    ak_array = ak.from_numpy(array)
    assert ak.to_list(ak.sum(ak_array, axis=1)) == ak.to_list(array.sum(axis=1))

[agoose77]

@ianna did this segfault or raise an Exception? I see only a ValueError in master & 1.7.0 on my machine.

[ianna]

@ianna did this segfault or raise an Exception? I see only a ValueError in master & 1.7.0 on my machine.

yep, the first :-(

[ianna]

I think, I need some help here. It seems to me, I miss some extra information - starts? The starts can help in ordering the parents, but I'm not sure how to efficiently sort the distincts if the parents are ordered. Especially if the distincts have missing values (e.g. -1)..

This kernel calculates the distincts and does not assume that the parents are ordered. Then, this kernel calculates gaps (between the parents?) and either assumes(?) that the parents are ordered or is it intentionally(?) produces unordered gaps as a result? Finally, this kernel needs ordered gaps, I think.

[agoose77]

@ianna Yes, I'm pretty sure that findgaps and nonlocal_outstartstops are both assuming that the parents are ordered, because the < operator doesn't yield much meaningful information if this isn't true. Maybe a better way of putting it is that the currently behaviour is broken for non-monotonic increasing parents, so it reinforces my belief that these kernels assume monotonic increasing parents.

Please stop me if I am telling you things that you already know! This is just from my attempts to understand this logic:

The gaps is used to identify missing sublists in ordered parents (see my above comment). I haven't looked into why we don't just use parents (compute gaps directly in the kernel(s) that need it), but it is used by nonlocal_outstartstops to mark the break between sublists:
https://github.com/scikit-hep/awkward-1.0/blob/500e0dd06fa4aa542ac0226da24851fb730e5042/src/cpu-kernels/awkward_ListOffsetArray_reduce_nonlocal_outstartsstops_64.cpp#L32-L36
This is always at least one (each entry in gaps corresponds to a sublist), but can be more where we have gaps between parents (empty sublists)

However, I made an assumption about distincts that I think may not be true, and would explain why the fix doesn't work in some cases. Let me look.

[agoose77]

I'm looking at this further. Your test case is interesting @ianna , because it seems like the truly jagged case is not fixed by ordering the parents, which is why I missed it with the regular test case.

[agoose77]

How Jagged Reduction Works

This is for my benefit in referencing how we do jagged reduction.

There are two types of information running through jagged reduction:

• context for this content
• context for child content

Some of the context for this content is passed from the parent. In the case of the root content, the parent does not exist, so this information is given by the assumption of an outer list of length len(this).

I'm ignoring starts as it's only used by position-aware reducers.

From the parent (pertaining to this layout), we have:

• parents - the index of the parent sublist that each item in this list belongs to. For the top-level content,parents = [0, 0, ...] where len(parents) = len(content)
• outlength - the length of the return value at this level (1 for the root)

From this content, we compute the following for the reduction of our child:

• nextcarry - To ensure reducers act over adjacent values, we re-index the child content to guarantee this.
• nextparents - We want to ensure that the values to be reduced over are adjacent. nextparents is used to keep track of which reduction each child content item belongs to. I'll mention later why there are jumps in this sometimes.

From this content, we compute the following for our wrapped reduction result:

• outstarts - The start indices into the reduction result of our sublists
• outstops - The stop indices into the reduction result of our sublists
• gaps - An array of length outlength
• distincts - An array of length N (given by outlength * max sublist length M). I'm still figuring this out.

depth==1

For a list:

<ListOffsetArray len='4'>
    <offsets><Index dtype='int64' len='5'>[0 2 4 6 8]</Index></offsets>
    <content><ListOffsetArray len='8'>
        <offsets><Index dtype='int64' len='9'>[ 0  1  3  4  6  7  9 10 12]</Index></offsets>
        <content><NumpyArray dtype='int64' len='12'>
            [ 1  2  3  5  7 11 13 17 19 23 29 31]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>

    [
        # axis 0
        [
            # axis 1
            [1], [2, 3]
        ],
        [
            [5], [7, 11]
        ],
        [
            [13], [17, 19]
        ],
        [
            [23], [29, 31]
        ]
    ]

we have two items. As this is the outermost list, they all have the same parent (parents = [0, 0, 0, 0]).

If we want to compute np.prod(array, axis=0), then we want to take the product over each column, i.e. [1], [5], [13], [23].
Clearly we place our first column sublist indices together (0, 2, 4, 6) and give them the same parent (0, 0, 0, 0). The same holds for the second column, i.e. we go from:

[
      [1], [2, 3]
      [5], [7, 11]
      [13], [17, 19]
      [23], [29, 31]
]

to

[
    # 0    2     4      6                         (carry)
      [1], [5], [13], [23], # parent 0
    #   1         3        5        7       (carry)
      [2, 3], [7, 11], [17, 19], [20, 31] # parent 1
]

So, we have

• nextcarry = [0, 2, 4, 6, 8, 1, 3, 5, 7] - group columns together (1, 5, 13, ...) and ((2, 3), (7, 11), ...)
• nextparents = [0, 0, 0, 0, 1, 1, 1, 1] (see above)
  and
• distincts = [0, 0] ??
• gaps = [1] ??

We can recurse into the child with this same procedure: We first pass in nextparents and nextstarts to the child content (another list!)

depth==2

So, we have this list:

<ListOffsetArray len='8'>
    <offsets><Index dtype='int64' len='9'>[ 0  1  2  3  4  6  8 10 12]</Index></offsets>
    <content><NumpyArray dtype='int64' len='12'>
        [ 1  5 13 23  2  3  7 11 17 19 29 31]
    </NumpyArray></content>
</ListOffsetArray>

[
      [1], [5], [13], [23], # parent 0
      [2, 3], [7, 11], [17, 19], [20, 31] # parent 1
]

With

• parents = [0, 0, 0, 0, 1, 1, 1, 1]

In order to maintain the adjacency of reduced elements, we have to split our sublists to pull out the columns for the next reduction, i.e.

[
      1, 5, 13, 23, # column 0
      2, 7, 17, 29, # column 1
      3, 11, 19, 31 # column 2
]

Again, this is because once the reducer acts on adjacent elements, we will be left with three distinct values which correspond to the reduction result.

To obtain this ordering, we need a carry:

• nextcarry = [0, 1, 2, 3, 4, 6, 8, 10, 5, 7, 9, 11]
  And these sublists have unique parents
• nextparents = [0, 0, 0, 0, 2, 2, 2, 2, 3, 3, 3, 3]

So, why is there a jump in nextparents between 0 and 2? Well, the reducer behaviour for jagged arrays is to treat all sublists as having len(sublist)=maxlen. This is visually represented as though this current array were:

[
    [1, Ø], [5, Ø], [13, Ø], [23, Ø],
    [2, 3], [7, 11], [17, 19], [29, 31],
]

This would then look like:

[
      1, 5, 13, 23, # column 0
      Ø, Ø, Ø, Ø  # column 1 (missing)
      2, 7, 17, 29, # column 2
      3, 11, 19, 31 # column 3
]

i.e. we behave as though there is a second column between (1,5,13,23) and (2, 7, 17, 29) that needs to be accounted for

So we have

• distincts = [0, -1, 1, 0] - ??
• gaps = [1, 1] - ??

depth == 3

Finally, we have a NumpyArray with the following contents after carry'ing:

[ 1,  5, 13, 23,  2,  7, 17, 29,  3, 11, 19, 31]

And the reducer simply walks over the array and reduces within the same sublists (using parents).
We're then left with

[ 1495 , 1, 6902, 19437]