`ak.any` yields True for empty rows in float arrays with the jax backend

[ikrommyd]

I'm actually not sure if this is a bug. I actually don't know if ak.any should work at all in such cases. I'm just noticing and inconsistent behavior between the cpu and jax backends even though ideally ak.any should be called on boolean arrays.

To reproduce

In [1]: import awkward as ak

In [2]: ak.jax.register_and_check()

In [3]: ak.any(ak.Array([[1], [1, 2], [], [3, 4], [7]], backend="cpu"), axis=1)
Out[3]: <Array [True, True, False, True, True] type='5 * bool'>

In [4]: ak.any(ak.Array([[1], [1, 2], [], [3, 4], [7]], backend="jax"), axis=1)
WARNING:2025-04-13 19:20:37,179:jax._src.xla_bridge:967: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
Out[4]: <Array [True, True, True, True, True] type='5 * bool'>

It also yields incorrect length

In [1]: import awkward as ak

In [2]: ak.jax.register_and_check()

In [3]: ak.any(ak.Array([[], [1, 2], [], [3, 4], []], backend="cpu"), axis=1)
Out[3]: <Array [False, True, False, True, False] type='5 * bool'>

In [4]: ak.any(ak.Array([[], [1, 2], [], [3, 4], []], backend="jax"), axis=1)
WARNING:2025-04-13 19:17:46,219:jax._src.xla_bridge:967: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
Out[4]: <Array [True, True, True, True] type='4 * bool'>

but the current status of #3464 fixes that and only gives us incorrect values:

In [1]: import awkward as ak

In [2]: ak.jax.register_and_check()

In [3]: ak.any(ak.Array([[], [1, 2], [], [3, 4], []], backend="cpu"), axis=1)
Out[3]: <Array [False, True, False, True, False] type='5 * bool'>

In [4]: ak.any(ak.Array([[], [1, 2], [], [3, 4], []], backend="jax"), axis=1)
WARNING:2025-04-13 19:21:42,367:jax._src.xla_bridge:967: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
Out[4]: <Array [True, True, True, True, True] type='5 * bool'>

[ikrommyd]

Well on the other hand ak.all() on empty elements gives True

In [1]: import awkward as ak

In [2]: ak.all(ak.Array([[], [1., 2], [], [3, 4]], backend="cpu"), axis=1)
Out[2]: <Array [True, True, True, True] type='4 * bool'>

Can someone tell me what the original intention was here?

[ariostas]

Well on the other hand ak.all() on empty elements gives True

Can someone tell me what the original intention was here?

$(\forall s\in S\ : p(s))$ is vacuously true for any proposition $p$, if $S=\emptyset$. So ak.all should return true for any empty lists. (so it's working well)

$(\exists s\in S\ : p(s))$ is trivially false for any proposition $p$, if $S=\emptyset$. So ak.any should return false for any empty lists. (so the jax backend is indeed the wrong one)