Four-vector addition with jax backend

[alexander-held]

Version of Awkward Array

2.3.0

Description and code to reproduce

This came out of a pyhep.dev hacking session. We tried to build a simple analysis example: read four-vector information for jets, "calibrate" the jet energies (as a function of a nuisance parameter) and then take derivatives of the mean of invariant masses wrt. the nuisance parameter. A reproducer is provided at https://gist.github.com/alexander-held/9ca8eae6ab99572d7a332ab8c1b02beb.

Switching to the jax backend (without taking gradients) triggers

TypeError: Unexpected input type for array: <class 'awkward.contents.numpyarray.NumpyArray'>

This error occurred while calling

    numpy.add.__call__(
        <MomentumArray4D [{eta: -3.1967773, ...}, ..., {...}] type='140 * M...'>
        <MomentumArray4D [{eta: 1.2490234, ...}, ..., {...}] type='140 * Mo...'>
    )

at the point of four-vector addition. Changing this line to instead return the invariant mass of the first jet triggers another error:

RuntimeError: Cannot differentiate through count_zero

This error occurred while calling

    ak.mean(
        <Array [3.3886719, 6.2890625, ..., 4.1523438] type='140 * float32'>
        axis = None
        keepdims = False
    )

[agoose77]

Fixed by #2592