`ak.argmin/argmax` give incorrect results in the jax backend

[ikrommyd]

In [3]: import awkward as ak
   ...: ak.jax.register_and_check()
   ...:
   ...: x = ak.argmin(ak.Array([[1, 2], [], [6,4,5], [7]], backend="cpu"), axis=1)
   ...: y = ak.argmin(ak.Array([[1, 2], [], [6,4,5], [7]], backend="jax"), axis=1)

In [4]: x
Out[4]: <Array [0, None, 1, 0] type='4 * ?int64'>

In [5]: y
Out[5]: <Array [0, None, 3, 5] type='4 * ?int64'>

My understand is that this happens because positional corrections aren't applied like the cpu backend does:

awkward/src/awkward/_reducers.py

Line 186 in c93da2f

apply_positional_corrections(result_array, parents, starts, shifts)

even though the output of the reduce kernel is correct.

[ianna]

@ikrommyd - thanks for opening the issue. The reducers do not support nans nor missing values.

[ikrommyd]

@ikrommyd - thanks for opening the issue. The reducers do not support nans nor missing values.

In [5]: import awkward as ak
   ...: ak.jax.register_and_check()
   ...: x = ak.argmin(ak.Array([[1, 2], [3], [6,4,5], [7]], backend="cpu"), axis=1)
   ...: y = ak.argmin(ak.Array([[1, 2], [3], [6,4,5], [7]], backend="jax"), axis=1)

In [6]: x
Out[6]: <Array [0, 0, 1, 0] type='4 * ?int64'>

In [7]: y
Out[7]: <Array [0, 2, 4, 6] type='4 * ?int64'>

NaNs or missing values don't play a role here. It's the actual values that are wrong here. argmin spits back incorrect values that have no meaning. The only problem is that positional corrections are not applied for those functions.