Skip to content

feat: add argmax implementation using CCCL#3777

Merged
ianna merged 21 commits intoscikit-hep:mainfrom
maxymnaumchyk:maxymnaumchyk/cccl-kernel
Jan 12, 2026
Merged

feat: add argmax implementation using CCCL#3777
ianna merged 21 commits intoscikit-hep:mainfrom
maxymnaumchyk:maxymnaumchyk/cccl-kernel

Conversation

@maxymnaumchyk
Copy link
Copy Markdown
Collaborator

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
Copy link
Copy Markdown
Collaborator Author

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
Copy link
Copy Markdown

codecov bot commented Dec 19, 2025

Codecov Report

❌ Patch coverage is 0% with 41 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.54%. Comparing base (c6fc079) to head (0058c5f).
⚠️ Report is 3 commits behind head on main.

Files with missing lines Patch % Lines
src/awkward/_connect/cuda/_compute.py 0.00% 38 Missing ⚠️
src/awkward/_backends/cupy.py 0.00% 3 Missing ⚠️
Additional details and impacted files
Files with missing lines Coverage Δ
src/awkward/_connect/cuda/__init__.py 0.00% <ø> (ø)
src/awkward/_backends/cupy.py 45.45% <0.00%> (-2.17%) ⬇️
src/awkward/_connect/cuda/_compute.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
Copy link
Copy Markdown

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

@maxymnaumchyk maxymnaumchyk marked this pull request as ready for review January 12, 2026 11:56
Copy link
Copy Markdown
Member

@ianna ianna left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@maxymnaumchyk - This is great! Please go ahead and merge it if you are done with it. Thanks!

Co-authored-by: Ianna Osborne <ianna.osborne@cern.ch>
@ianna ianna merged commit d47fbef into scikit-hep:main Jan 12, 2026
65 of 67 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants