feat: support record reducer overloads

[agoose77]

Fixes #1423

• Implement non-positional reducer support
• Implement positional reducer support
• Changed the ak.behavior documentation to use MyST Markdown

[agoose77]

This seems simple enough; build a list over the RecordArray which the custom reducer reduces along axis=-1. However, the open question is how to handle masking/identity. The simplest interface is to always require the custom reducer to include an identity element for empty sublists. This can then be masked out by the dispatcher. We currently do this for our normal reduction pathways.

However, doing this with existing highlevel primitives may not be trivial. e.g. an implementation of ak.max for a record structure, which looks at the first slot:

def _my_record_max(array):
    key = array["0"]
    j = ak.from_regular(
        ak.argmax(key, keepdims=True, mask_identity=False, axis=-1)
    )
    return array[j][..., 0]


ak.behavior[(ak.max, "my_record")] = _tuples_max

This will fail if we have an empty sublist and we don't mask the reducer.

One of the constraints that we are operating under is the desire not to have third party code using our nplike abstraction. Whilst for NumPy, CuPy, and JAX that should be possible (just type-check the buffers), typetracer would prove more fiddly (incidentally, this does add a vote in favour of making our typetracer nplike L2).

The above can be extended to support empty sublists:

def _option_is_trivial(array):
    any_is_none = ak.any(ak.is_none(array, axis=0), axis=None)
    return not (ak.typetracer.is_unknown_scalar(any_is_none) or any_is_none)


def _my_record_max(array):
    key = array["0"]
    j = ak.from_regular(
        ak.argmax(key, keepdims=True, mask_identity=True, axis=-1)
    )
    out = ak.to_layout(
        array[j][..., 0]
    )
    assert out.is_option
    
    if _option_is_trivial(out):
        out = out.to_IndexedOptionArray64()
        return ak.contents.IndexedArray(
            out.index,
            out.content
        )
    else:
        return ak.fill_none(out, identity_element)

but where there are any empty sublists, it will involve a copy of out in ak.fill_none even if the output is subsequently masked.

I think the proper API therefore includes the mask_identity flag. Here's an attempt to write a simple max that doesn't copy the RecordArray unless it has no choice, without using nplike.

def _option_is_trivial(array):
    any_is_none = ak.any(ak.is_none(array, axis=0), axis=None)
    return not (ak.typetracer.is_unknown_scalar(any_is_none) or any_is_none)


def _my_record_max(array: ak.Array, mask_identity: bool):
    key = array["0"]
    j = ak.from_regular(ak.argmax(key, keepdims=True, mask_identity=True, axis=-1))
    out = ak.to_layout(
        array[j][..., 0]
    )

    if mask_identity:
        return out
    # Avoid content `_carry` for options with empty masks
    elif _option_is_trivial(out):
        out = out.to_IndexedOptionArray64()
        return ak.contents.IndexedArray(
            out.index,
            out.content
        )
    # _carry the content (and fill missing values)
    else:
        return ak.fill_none(out, identity_element)

Meanwhile, argmax is trivial!

def _my_record_argmax(array: ak.Array, mask_identity: bool):
    return ak.argmax(array["0"], keepdims=False, mask_identity=False, axis=-1)

[codecov]

Codecov Report

Merging #2458 (872de48) into main (4468cf5) will decrease coverage by 0.01%.
The diff coverage is 75.86%.

Additional details and impacted files

Impacted Files	Coverage Δ
src/awkward/operations/ak_with_name.py	100.00% <ø> (ø)
src/awkward/_do.py	83.22% <50.00%> (ø)
src/awkward/contents/recordarray.py	84.68% <77.77%> (-0.44%)	⬇️

[jpivarski]

This is a good argument for the API requiring users to write functions that include a mask_identity argument. (Since these are functions users write, it's much harder to add required arguments later!) I don't suppose it needs a keepdims argument, since keeping a dimension should always be done in a regular way—behavior developers shouldn't have control over that.

In principle, the behavior developers might want to raise an exception on mask_identity=False.

Although you have unit tests for the sum-of-vectors case, it could be good to test-implement it in Vector (here) to be sure that the API is what we need. (Notice all of the __numba_typer__/__numba_lower__ behaviors defined in this file: the API in Awkward and its use in Vector were developed in tandem.)

[agoose77]

I agree. In truth, I'm currently thinking about how vector should support this - the dispatch machinery is fairly long. I suspect it will just require simply writing the same code!

[agoose77]

@jpivarski inspired by the existing (and more complex) binary addition, I took a stab at sum here.

For vector, I can't initially think of any reducers that we'd want to implement that do not have identities. Can you think of any?

In the event that there are none, the idea here is that we can re-use the sum implementation for numpy-backed vectors too (by moveaxis, reduce, moveaxis, to prepare the array for reduction)

[agoose77]

My implementation here implicitly assumes that records are atoms, i.e. one cannot reduce deeper than the record itself. This was motivated by RecordArray.purelist_depth which returns 1, and justifies not passing axis to the custom reducer. I believe this holds, but weigh in if you think that's a misstep.

[jpivarski]

For vector, I can't initially think of any reducers that we'd want to implement that do not have identities. Can you think of any?

In Vector (and elsewhere), the only reducer I can think of wanting to override at all is sum. The any/all pair applies to booleans, which are not geometrical vectors, the min/max pair applies to objects with an ordering, and $n>1$-dimensional vectors do not have an ordering; same for argmin/argmax. I can't even see a reason to have prod (scalar product is not reducible with more than two arguments; vector product only has a definition for 3D vectors; and neither of these should be confused with *, which only combines a scalar with a vector).

Two reducers that can be implemented for Vectors: count and count_nonzero (where "nonzero" means "nonzero magnitude"). They have identities, though.

My implementation here implicitly assumes that records are atoms, i.e. one cannot reduce deeper than the record itself.

That's right. We can enshrine that as a rule.

[agoose77]

I don't yet know where to put this. It feels a bit odd creating a highlevel array inside code written in a Content class definition, so I've moved it to a top-level function, much like the other reducer overloads are defined.