feat: add argmax implementation using CCCL

[maxymnaumchyk]

This was working on a simple test like this one just yesterday:

import awkward as ak
awkward_array = ak.Array([[1], [2, 3], [4, 5, 6, 7, 8, 9], [], [], [10]], backend = 'cuda')
result = ak.argmax(awkward_array, axis = 1)

Now I'm getting a NumbaNotImplementedError: input_data.dtype error for some reason. I'll be trying to fix this ~

[maxymnaumchyk]

Full error:

---------------------------------------------------------------------------
NumbaNotImplementedError                  Traceback (most recent call last)
Cell In[4], line 1
----> 1 result = ak.argmax(awkward_array, axis = 1)

File ~/awkward/src/awkward/_dispatch.py:41, in named_high_level_function.<locals>.dispatch(*args, **kwargs)
     38 @wraps(func)
     39 def dispatch(*args, **kwargs):
     40     # NOTE: this decorator assumes that the operation is exposed under `ak.`
---> 41     with OperationErrorContext(name, args, kwargs):
     42         gen_or_result = func(*args, **kwargs)
     43         if isgenerator(gen_or_result):

File ~/awkward/src/awkward/_errors.py:80, in ErrorContext.__exit__(self, exception_type, exception_value, traceback)
     78     self._slate.__dict__.clear()
     79     # Handle caught exception
---> 80     raise self.decorate_exception(exception_type, exception_value)
     81 else:
     82     # Step out of the way so that another ErrorContext can become primary.
     83     if self.primary() is self:

File ~/awkward/src/awkward/_dispatch.py:67, in named_high_level_function.<locals>.dispatch(*args, **kwargs)
     65 # Failed to find a custom overload, so resume the original function
     66 try:
---> 67     next(gen_or_result)
     68 except StopIteration as err:
     69     return err.value

File ~/awkward/src/awkward/operations/ak_argmax.py:79, in argmax(array, axis, keepdims, mask_identity, highlevel, behavior, attrs)
     76 yield (array,)
     78 # Implementation
---> 79 return _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs)

File ~/awkward/src/awkward/operations/ak_argmax.py:163, in _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs)
    159 axis = regularize_axis(axis, none_allowed=True)
    161 reducer = ak._reducers.ArgMax()
--> 163 out = ak._do.reduce(
    164     layout,
    165     reducer,
    166     axis=axis,
    167     mask=mask_identity,
    168     keepdims=keepdims,
    169     behavior=ctx.behavior,
    170 )
    172 wrapped_out = ctx.wrap(
    173     out,
    174     highlevel=highlevel,
    175     allow_other=True,
    176 )
    178 # propagate named axis to output

File ~/awkward/src/awkward/_do.py:321, in reduce(layout, reducer, axis, mask, keepdims, behavior)
    319 parents = ak.index.Index64.zeros(layout.length, layout.backend.nplike)
    320 shifts = None
--> 321 next = layout._reduce_next(
    322     reducer,
    323     negaxis,
    324     starts,
    325     shifts,
    326     parents,
    327     1,
    328     mask,
    329     keepdims,
    330     behavior,
    331 )
    333 return next[0]

File ~/awkward/src/awkward/contents/listoffsetarray.py:1623, in ListOffsetArray._reduce_next(self, reducer, negaxis, starts, shifts, parents, outlength, mask, keepdims, behavior)
   1620 trimmed = self._content[self.offsets[0] : self.offsets[-1]]
   1621 nextstarts = self.offsets[:-1]
-> 1623 outcontent = trimmed._reduce_next(
   1624     reducer,
   1625     negaxis,
   1626     nextstarts,
   1627     shifts,
   1628     nextparents,
   1629     globalstarts_length,
   1630     mask,
   1631     keepdims,
   1632     behavior,
   1633 )
   1635 outoffsets = Index64.empty(outlength + 1, nplike)
   1636 assert outoffsets.nplike is nplike and parents.nplike is nplike

File ~/awkward/src/awkward/contents/numpyarray.py:1142, in NumpyArray._reduce_next(self, reducer, negaxis, starts, shifts, parents, outlength, mask, keepdims, behavior)
   1139 assert self.is_contiguous
   1140 assert self._data.ndim == 1
-> 1142 out = reducer.apply(self, parents, starts, shifts, outlength)
   1144 if mask:
   1145     outmask = ak.index.Index8.empty(outlength, self._backend.nplike)

File ~/awkward/src/awkward/_reducers.py:239, in ArgMax.apply(self, array, parents, starts, shifts, outlength)
    236 else:
    237     assert parents.nplike is array.backend.nplike
    238     array.backend.maybe_kernel_error(
--> 239         array.backend[
    240             "awkward_reduce_argmax",
    241             result.dtype.type,
    242             kernel_array_data.dtype.type,
    243             parents.dtype.type,
    244         ](
    245             result,
    246             kernel_array_data,
    247             parents.data,
    248             parents.length,
    249             outlength,
    250         )
    251     )
    252 result_array = ak.contents.NumpyArray(result, backend=array.backend)
    253 print("We are wrapping up with contents!!", result_array)

File ~/awkward/src/awkward/_kernels.py:239, in CudaComputeKernel.__call__(self, *args)
    237 print("with", len(args), "parameters.")
    238 print("Parameters:", args, "\n")
--> 239 return self._impl(*args)

File ~/awkward/src/awkward/_connect/cuda/cccl_kernels/awkward_reduce_argmax.py:54, in awkward_reduce_argmax(result, input_data, parents_data, parents_length, outlength)
     51 print(type(input_data))
     52 print(type(parents_data))
---> 54 @gpu_struct
     55 class ak_array:
     56     data: input_data.dtype
     57     local_index: cp.int64

File ~/miniforge3/envs/ak3/lib/python3.13/site-packages/cuda/compute/struct.py:165, in gpu_struct(field_dict, name)
    162 assert isinstance(field_dict, dict)
    164 # Convert all fields to numba types
--> 165 field_spec = {key: _convert_field_type(val) for key, val in field_dict.items()}
    167 # Create base struct
    168 struct_class = make_struct_type(
    169     name,
    170     tuple(field_spec.keys()),
    171     tuple(field_spec.values()),
    172 )

File ~/miniforge3/envs/ak3/lib/python3.13/site-packages/cuda/compute/struct.py:515, in _convert_field_type(val)
    513 if _is_struct_type(val):
    514     return val
--> 515 return numba.np.numpy_support.from_dtype(val)

File ~/miniforge3/envs/ak3/lib/python3.13/site-packages/numba/np/numpy_support.py:137, in from_dtype(dtype)
    134         subtype = from_dtype(dtype.subdtype[0])
    135         return types.NestedArray(subtype, dtype.shape)
--> 137 raise errors.NumbaNotImplementedError(dtype)

NumbaNotImplementedError: input_data.dtype

This error occurred while calling

    ak.argmax(
        <Array [[1], [2, 3], [4, ..., 9], ..., [], [10]] type='6 * var * int64'>
        axis = 1
    )

[codecov]

Codecov Report

❌ Patch coverage is 8.13953% with 79 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.54%. Comparing base (d980f98) to head (bd6ac7d).

Files with missing lines	Patch %	Lines
...connect/cuda/cccl_kernels/awkward_reduce_argmax.py	0.00%	54 Missing ⚠️
src/awkward/_backends/cupy.py	20.00%	12 Missing ⚠️
src/awkward/_kernels.py	21.42%	11 Missing ⚠️
src/awkward/_connect/cuda/cccl_kernels/__init__.py	0.00%	2 Missing ⚠️

Additional details and impacted files

Files with missing lines	Coverage Δ
src/awkward/_reducers.py	98.12% <100.00%> (+<0.01%)	⬆️
src/awkward/_connect/cuda/cccl_kernels/__init__.py	0.00% <0.00%> (ø)
src/awkward/_kernels.py	66.43% <21.42%> (-4.78%)	⬇️
src/awkward/_backends/cupy.py	47.50% <20.00%> (-19.17%)	⬇️
...connect/cuda/cccl_kernels/awkward_reduce_argmax.py	0.00% <0.00%> (ø)

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

The documentation preview is ready to be viewed at http://preview.awkward-array.org.s3-website.us-east-1.amazonaws.com/PR3777