awkward icon indicating copy to clipboard operation
awkward copied to clipboard

`IndexError` in `ak.combinations()` with `cuda` backend

Open kmohrman opened this issue 7 months ago • 3 comments

Version of Awkward Array

2.8.3

Description and code to reproduce

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

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

print(ak.__version__)

def test(pt,eta,phi,mass,backend_type):

    Jet_pt = ak.Array(pt,backend=backend_type)
    Jet_eta = ak.Array(eta,backend=backend_type)
    Jet_phi = ak.Array(phi,backend=backend_type)
    Jet_mass = ak.Array(mass,backend=backend_type)

    jets = ak.zip(
        {
            "pt": Jet_pt,
            "eta": Jet_eta,
            "phi": Jet_phi,
            "mass": Jet_mass,
        },
        with_name="PtEtaPhiMLorentzVector",
        behavior=candidate.behavior,
    )

    print("\njets backend",ak.backend(jets))
    jet_combs = ak.combinations(jets, 3, fields=["j1", "j2", "j3"])
    result = jet_combs.j1 + jet_combs.j2 + jet_combs.j3


pt   = [[71.74248, 27.820534], [23.159304, 21.095522, 17.393976]]
eta  = [[-1.6034726, -0.7512972], [-2.314745, 3.4676905, -0.5865479]]
phi  = [[-2.251697, 0.56961906], [0.21612246, -0.79849535, -2.8226633]]
mass = [[12.71695, 7.453251], [4.6756053, 5.584647, 3.4867086]]

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

The jet_combs = ak.combinations(jets, 3, fields=["j1", "j2", "j3"]) line runs fine with cpu, but with cuda, the same line gives an error IndexError: cannot slice RecordArray (of length 5) with array([99]).

The full traceback is pasted here:

Traceback (most recent call last):
  File "/home/k.mohrman/coffea_dir/gpu_studies/columnar_gpu/t.py", line 36, in <module>
    test(pt,eta,phi,mass,"cuda")
  File "/home/k.mohrman/coffea_dir/gpu_studies/columnar_gpu/t.py", line 26, in test
    jet_combs = ak.combinations(jets, 3, fields=["j1", "j2", "j3"])
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env5/lib/python3.12/site-packages/awkward/_dispatch.py", line 41, in dispatch
    with OperationErrorContext(name, args, kwargs):
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env5/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_env5/lib/python3.12/site-packages/awkward/_dispatch.py", line 67, in dispatch
    next(gen_or_result)
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env5/lib/python3.12/site-packages/awkward/operations/ak_combinations.py", line 195, in combinations
    return _impl(
           ^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env5/lib/python3.12/site-packages/awkward/operations/ak_combinations.py", line 238, in _impl
    out = ak._do.combinations(
          ^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env5/lib/python3.12/site-packages/awkward/_do.py", line 142, in combinations
    return layout._combinations(n, replacement, recordlookup, parameters, axis, 1)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env5/lib/python3.12/site-packages/awkward/contents/listoffsetarray.py", line 1446, in _combinations
    contents.append(self._content._carry(ptr, True))
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/blue/p.chang/k.mohrman/dir_for_miniconda/miniconda3/envs/coffeagpu_env5/lib/python3.12/site-packages/awkward/contents/recordarray.py", line 533, in _carry
    raise ak._errors.index_error(self, where)
IndexError: cannot slice RecordArray (of length 5) with array([99])

This error occurred while calling

    ak.combinations(
        <PtEtaPhiMLorentzVectorArray [[{pt: ..., ...}, ...], ...] type='2 *...'>
        3
        fields = ['j1', 'j2', 'j3']
    )

kmohrman avatar Jun 04 '25 02:06 kmohrman

@kmohrman - thanks for reporting it! Indeed, scalars are one dimensional arrays in cupy (unlike in numpy).

ianna avatar Jun 04 '25 11:06 ianna

@kmohrman - Thanks for reporting it! It looks like this test should have raised a "not implemented error": the cuda combinations kernel supports only n == 2 for now. I marked this issue as feature request.

ianna avatar Jun 24 '25 07:06 ianna

Perhaps this uses this kernel: https://github.com/scikit-hep/awkward/blob/473029b8cc7345083c29186e0c3def874eb3a9f2/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_combinations_length.cu which doesn't check that n==2 like here for example: https://github.com/scikit-hep/awkward/blob/473029b8cc7345083c29186e0c3def874eb3a9f2/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_combinations.cu#L47-L53 to raise a better error

ikrommyd avatar Jun 27 '25 17:06 ikrommyd