`ValueError` with `cuda` backend

[kmohrman]

Version of Awkward Array

2.8.1

Description and code to reproduce

For an operation that works fine for the cpu backend, I am getting a ValueError error when running with the cuda backend. This code reproduces the issue:

import awkward as ak
from coffea.nanoevents.methods import candidate

def test(pt,eta,phi,mass,backend_type):
    jets = ak.zip(
        {
            "pt"  : ak.Array(pt,   backend=backend_type),
            "eta" : ak.Array(eta,  backend=backend_type),
            "phi" : ak.Array(phi,  backend=backend_type),
            "mass": ak.Array(mass, backend=backend_type)
        },
        with_name="PtEtaPhiMLorentzVector",
        behavior=candidate.behavior,
    )

    jet_combs = ak.combinations(jets, 2, fields=["j1", "j2"])
    print("\njet_combs backend:",ak.backend(jet_combs))
    result = jet_combs.j1 + jet_combs.j2

pt   = [[23.1, 21.0, 17.3], [72.6, 51.3, 33.7]]
eta  = [[-2.3, 3.4, -0.5], [1.7, 2.2, 0.7]]
phi  = [[0.2, -0.7, -2.8], [-2.5, -0.1, 2.1]]
mass = [[4.6, 5.5, 3.4], [12.0, 9.5, 7.8]]

test(pt,eta,phi,mass,"cpu")
test(pt,eta,phi,mass,"cuda")

The result = jet_combs.j1 + jet_combs.j2 line runs fine with cpu, but with cuda, the same line gives an error ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all(). The full error is pasted here:

Traceback (most recent call last):
  File "/home/k.mohrman/coffea_dir/gpu_studies/columnar_gpu/t.py", line 28, in <module>
    test(pt,eta,phi,mass,"cuda")
  File "/home/k.mohrman/coffea_dir/gpu_studies/columnar_gpu/t.py", line 20, in test
    result = jet_combs.j1 + jet_combs.j2
             ~~~~~~~~~~~~~^~~~~~~~~~~~~~
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_operators.py", line 54, in func
    return ufunc(self, other)
           ^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/highlevel.py", line 1619, in __array_ufunc__
    with ak._errors.OperationErrorContext(name, inputs, kwargs):
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_errors.py", line 80, in __exit__
    raise self.decorate_exception(exception_type, exception_value)
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/highlevel.py", line 1620, in __array_ufunc__
    return ak._connect.numpy.array_ufunc(ufunc, method, inputs, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_connect/numpy.py", line 469, in array_ufunc
    out = ak._broadcasting.broadcast_and_apply(
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1200, in broadcast_and_apply
    out = apply_step(
          ^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1178, in apply_step
    return continuation()
           ^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1147, in continuation
    return broadcast_any_list()
           ^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 671, in broadcast_any_list
    outcontent = apply_step(
                 ^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1178, in apply_step
    return continuation()
           ^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1147, in continuation
    return broadcast_any_list()
           ^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 747, in broadcast_any_list
    outcontent = apply_step(
                 ^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1178, in apply_step
    return continuation()
           ^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1130, in continuation
    return broadcast_any_indexed()
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1080, in broadcast_any_indexed
    return apply_step(
           ^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1160, in apply_step
    result = action(
             ^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_connect/numpy.py", line 381, in action
    return _array_ufunc_adjust(custom, inputs, kwargs, behavior)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_connect/numpy.py", line 184, in _array_ufunc_adjust
    out = custom(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/vector/_methods.py", line 3997, in add
    return module.add.dispatch(self, other)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/vector/_compute/lorentz/add.py", line 207, in dispatch
    handler._wrap_dispatched_function(function)(
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/vector/backends/awkward.py", line 1093, in __call__
    return ak.transform(transformer, *awkward_arrays)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_dispatch.py", line 64, in dispatch
    next(gen_or_result)
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/operations/ak_transform.py", line 469, in transform
    return _impl(
           ^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/operations/ak_transform.py", line 632, in _impl
    out = apply_broadcasting_step(
          ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1178, in apply_step
    return continuation()
           ^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1147, in continuation
    return broadcast_any_list()
           ^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 671, in broadcast_any_list
    outcontent = apply_step(
                 ^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1178, in apply_step
    return continuation()
           ^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1130, in continuation
    return broadcast_any_indexed()
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1080, in broadcast_any_indexed
    return apply_step(
           ^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1160, in apply_step
    result = action(
             ^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/awkward/operations/ak_transform.py", line 592, in action
    out = transformation(tuple(inputs), **kwargs, behavior=behavior)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/vector/backends/awkward.py", line 1070, in transformer
    out_numpys = bind(self.func, *args2bind)(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/vector/backends/awkward.py", line 1015, in __call__
    return self.func(*args, *iargs, **keywords)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/vector/_compute/lorentz/add.py", line 148, in f
    coord1, coord2, coord3 = spatial_add(
                             ^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/vector/_compute/spatial/add.py", line 393, in rhophi_eta_rhophi_eta
    return (rho, phi, eta.rhophi_z(lib, rho, phi, z1 + z2))
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/vector/_compute/spatial/eta.py", line 56, in rhophi_z
    return lib.nan_to_num(
           ^^^^^^^^^^^^^^^
  File "cupy/_core/core.pyx", line 1766, in cupy._core.core._ndarray_base.__array_function__
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/cupy/_math/misc.py", line 422, in nan_to_num
    nan = _check_nan_inf(nan, dtype)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env1/lib/python3.12/site-packages/cupy/_math/misc.py", line 402, in _check_nan_inf
    elif cupy.isnan(x):
         ^^^^^^^^^^^^^
  File "cupy/_core/core.pyx", line 1293, in cupy._core.core._ndarray_base.__nonzero__
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

This error occurred while calling

    numpy.add.__call__(
        <PtEtaPhiMLorentzVectorArray [[{pt: 23.1, ...}, ...], ...] type='2 ...'>
        <PtEtaPhiMLorentzVectorArray [[{pt: 21.0, ...}, ...], ...] type='2 ...'>
    )

[ikrommyd]

This is not an awkward problem. This is A) a vector problem and B) a cupy problem. The nan argument of nan_to_num should be a float. Yet in vector we pass in nan=cupy.nan_to_num((z != 0) * inf, posinf=nan) which is an array so we shouldn't do that.

In [1]: import cupy

In [2]: from math import inf, nan

In [3]: rho = cupy.array([39.72022144,  6.45279762, 19.3253414 , 49.08952218, 76.53535633, 41.62411281])

In [4]: z = cupy.array([200.22881345, -123.05876649,  305.25768238,  420.72237365, 217.63714915, 254.213766])

In [5]: cupy.nan_to_num(
   ...:         cupy.arcsinh(z / rho),
   ...:         nan=cupy.nan_to_num((z != 0) * inf, posinf=nan),
   ...:         posinf=inf,
   ...:         neginf=-inf,
   ...:     )
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[5], line 1
----> 1 cupy.nan_to_num(
      2         cupy.arcsinh(z / rho),
      3         nan=cupy.nan_to_num((z != 0) * inf, posinf=nan),
      4         posinf=inf,
      5         neginf=-inf,
      6     )

File /opt/conda/lib/python3.10/site-packages/cupy/_math/misc.py:422, in nan_to_num(x, copy, nan, posinf, neginf)
    420     out = cupy.empty_like(x) if copy else x
    421 dtype = out.dtype
--> 422 nan = _check_nan_inf(nan, dtype)
    423 posinf = _check_nan_inf(posinf, dtype, False)
    424 neginf = _check_nan_inf(neginf, dtype, True)

File /opt/conda/lib/python3.10/site-packages/cupy/_math/misc.py:402, in _check_nan_inf(x, dtype, neg)
    400 elif x is None and neg is not None:
    401     x = cupy.finfo(dtype).min if neg else cupy.finfo(dtype).max
--> 402 elif cupy.isnan(x):
    403     x = cupy.nan
    404 elif cupy.isinf(x):

File cupy/_core/core.pyx:1293, in cupy._core.core._ndarray_base.__nonzero__()

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

However, numpy is fine with it:

In [6]: import numpy

In [7]: rho = numpy.array([39.72022144,  6.45279762, 19.3253414 , 49.08952218, 76.53535633, 41.62411281])

In [8]: z = numpy.array([200.22881345, -123.05876649,  305.25768238,  420.72237365, 217.63714915, 254.213766])

In [9]: numpy.nan_to_num(
   ...:         numpy.arcsinh(z / rho),
   ...:         nan=numpy.nan_to_num((z != 0) * inf, posinf=nan),
   ...:         posinf=inf,
   ...:         neginf=-inf,
   ...:     )
Out[9]: 
array([ 2.32044354, -3.64198212,  3.45388668,  2.844861  ,  1.76779802,
        2.50927908])

[ikrommyd]

@Saransh-cpp do you know why we pass an array here:

vector/src/vector/_compute/spatial/eta.py

Line 58 in 1513af6

nan=lib.nan_to_num((z != 0) * inf, posinf=nan),

?
There is also another case right above:

vector/src/vector/_compute/spatial/eta.py

Line 36 in 1513af6

nan=lib.nan_to_num((z != 0) * inf, posinf=nan),

and maybe more that I didn't notice using github search.

Docs say that it should be an int/float: https://numpy.org/doc/2.1/reference/generated/numpy.nan_to_num.html but numpy is fine with an array while cupy isn't. I'm leaning to believe that this is a numpy inconsistency and not a cupy one actually. Numpy is doing more than promised here.

[ikrommyd]

The inconsistency is reported under cupy: cupy/cupy#9143 and numpy: numpy/numpy#29076

In any case, we shouldn't be trying to pass an array there in vector.

[kmohrman]

Hi @ikrommyd, apologies for reporting in the wrong place, and thank you very much for the followups.

[Saransh-cpp]

nan_to_num is not standardised by the array API standard (RFC: add nan_to_num to the array API spec), so CuPy and NumPy can freely have different behaviors, but I agree, NumPy's documentation is wrong here.

For using nan=lib.nan_to_num(...) in vector - I remember this was done when I had just started working on vector - #172. Array API was still being developed at that point and vector's SymPy backend was just an idea. Both of them are in production now, so I'm sure we can look at possible solutions again.

I'll take a look at this. Moving it to vector for now as it is indeed a vector issue.

[Saransh-cpp]

Also, FYI, ak.nan_to_num(...) supports nan as an array - https://awkward-array.org/doc/main/reference/generated/ak.nan_to_num.html#ak.nan_to_num (it mirrors the np.nan_to_num behavior).

[ikrommyd]

Hi @ikrommyd, apologies for reporting in the wrong place, and thank you very much for the followups.

That's totally fine, it's not your job to find where the error is internally.

[ikrommyd]

Also, FYI, ak.nan_to_num(...) supports nan as an array - https://awkward-array.org/doc/main/reference/generated/ak.nan_to_num.html#ak.nan_to_num (it mirrors the np.nan_to_num behavior).

Yeah that uses numpy (backend.nplike.nan_to_num) internally

In [4]: np.nan_to_num([np.nan, np.nan], nan=[1,2])
Out[4]: array([1., 2.])

[ikrommyd]

Also, FYI, ak.nan_to_num(...) supports nan as an array - https://awkward-array.org/doc/main/reference/generated/ak.nan_to_num.html#ak.nan_to_num (it mirrors the np.nan_to_num behavior).

Oh now my brain realized that due to this, ak.nan_to_num also doesn't work with the cuda backend if nan is not a scalar haha. CPU backend is of course fine.

In [13]: x = ak.Array([np.nan, np.nan], backend="cuda")

In [14]: y = ak.Array([1, 2], backend="cuda")

In [15]: ak.nan_to_num(x, nan=y)

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

This error occurred while calling

    ak.nan_to_num(
        <Array [nan, nan] type='2 * float64'>
        nan = <Array [1, 2] type='2 * int64'>
    )

[Saransh-cpp]

Ah yes, I didn't realise that either. Maybe this should be an issue in awkward as well (indicate in docs for now at least).