Skip to content

"cuda" target in numba.vectorize not working correctly? #3179

@lgray

Description

@lgray

Version of Awkward Array

2.6.6

Description and code to reproduce

The following code fails:

import awkward as ak import cupy as cp import numba as nb ak.numba.register_and_check() @nb.vectorize(  [  nb.float32(nb.float32),  nb.float64(nb.float64),  ] ) def _square(x): return x * x @nb.vectorize(  [  nb.float32(nb.float32),  nb.float64(nb.float64),  ],  target="cuda", ) def _square_cuda(x): return x * x counts = cp.random.poisson(lam=3, size=50) flat_values = cp.random.normal(size=int(counts.sum())) values = ak.unflatten(flat_values, counts) values2_cpu = _square(ak.to_backend(values, "cpu")) print(values2_cpu) values2 = _square_cuda(values) print(values2)

resulting in the output:

[[0.045, 1.83], [1.55, 0.224, 0.621, ..., 1.24, 3.87], ..., [0.153, 0.0017]] --------------------------------------------------------------------------- ValueError Traceback (most recent call last) Cell In[12], line 33 29 values2_cpu = _square(ak.to_backend(values, "cpu")) 31 print(values2_cpu) ---> 33 values2 = _square_cuda(values) 35 print(values2) File ~/.conda/envs/coffea-gpu/lib/python3.11/site-packages/numba/cuda/vectorizers.py:28, in CUDAUFuncDispatcher.__call__(self, *args, **kws) 17 def __call__(self, *args, **kws): 18 """ 19 *args: numpy arrays or DeviceArrayBase (created by cuda.to_device). 20 Cannot mix the two types in one call. (...) 26 the input arguments. 27 """ ---> 28 return CUDAUFuncMechanism.call(self.functions, args, kws) File ~/.conda/envs/coffea-gpu/lib/python3.11/site-packages/numba/np/ufunc/deviceufunc.py:254, in UFuncMechanism.call(cls, typemap, args, kws) 252 # Begin call resolution 253 cr = cls(typemap, args) --> 254 args = cr.get_arguments() 255 resty, func = cr.get_function() 257 outshape = args[0].shape File ~/.conda/envs/coffea-gpu/lib/python3.11/site-packages/numba/np/ufunc/deviceufunc.py:202, in UFuncMechanism.get_arguments(self) 198 def get_arguments(self): 199 """Prepare and return the arguments for the ufunc. 200 Does not call to_device(). 201 """ --> 202 self._fill_arrays() 203 self._fill_argtypes() 204 self._resolve_signature() File ~/.conda/envs/coffea-gpu/lib/python3.11/site-packages/numba/np/ufunc/deviceufunc.py:100, in UFuncMechanism._fill_arrays(self) 98 self.scalarpos.append(i) 99 else: --> 100 self.arrays[i] = np.asarray(arg) File [~/coffea-gpu/awkward/src/awkward/highlevel.py:1434](https://analytics-hub.fnal.gov/user/lagray/lab/tree/coffea-gpu/coffea-gpu/awkward/src/awkward/highlevel.py#line=1433), in Array.__array__(self, dtype) 1429 with ak._errors.OperationErrorContext( 1430 "numpy.asarray", (self,), {"dtype": dtype} 1431 ): 1432 from awkward._connect.numpy import convert_to_array -> 1434 return convert_to_array(self._layout, dtype=dtype) File [~/coffea-gpu/awkward/src/awkward/_connect/numpy.py:481](https://analytics-hub.fnal.gov/user/lagray/lab/tree/coffea-gpu/coffea-gpu/awkward/src/awkward/_connect/numpy.py#line=480), in convert_to_array(layout, dtype) 480 def convert_to_array(layout, dtype=None): --> 481 out = ak.operations.to_numpy(layout, allow_missing=False) 482 if dtype is None: 483 return out File [~/coffea-gpu/awkward/src/awkward/_dispatch.py:64](https://analytics-hub.fnal.gov/user/lagray/lab/tree/coffea-gpu/coffea-gpu/awkward/src/awkward/_dispatch.py#line=63), in named_high_level_function.<locals>.dispatch(*args, **kwargs) 62 # Failed to find a custom overload, so resume the original function 63 try: ---> 64 next(gen_or_result) 65 except StopIteration as err: 66 return err.value File [~/coffea-gpu/awkward/src/awkward/operations/ak_to_numpy.py:48](https://analytics-hub.fnal.gov/user/lagray/lab/tree/coffea-gpu/coffea-gpu/awkward/src/awkward/operations/ak_to_numpy.py#line=47), in to_numpy(array, allow_missing) 45 yield (array,) 47 # Implementation ---> 48 return _impl(array, allow_missing) File [~/coffea-gpu/awkward/src/awkward/operations/ak_to_numpy.py:60](https://analytics-hub.fnal.gov/user/lagray/lab/tree/coffea-gpu/coffea-gpu/awkward/src/awkward/operations/ak_to_numpy.py#line=59), in _impl(array, allow_missing) 57 backend = NumpyBackend.instance() 58 numpy_layout = layout.to_backend(backend) ---> 60 return numpy_layout.to_backend_array(allow_missing=allow_missing) File [~/coffea-gpu/awkward/src/awkward/contents/content.py:1020](https://analytics-hub.fnal.gov/user/lagray/lab/tree/coffea-gpu/coffea-gpu/awkward/src/awkward/contents/content.py#line=1019), in Content.to_backend_array(self, allow_missing, backend) 1018 else: 1019 backend = regularize_backend(backend) -> 1020 return self._to_backend_array(allow_missing, backend) File [~/coffea-gpu/awkward/src/awkward/contents/listoffsetarray.py:2072](https://analytics-hub.fnal.gov/user/lagray/lab/tree/coffea-gpu/coffea-gpu/awkward/src/awkward/contents/listoffsetarray.py#line=2071), in ListOffsetArray._to_backend_array(self, allow_missing, backend) 2070 return buffer.view(np.dtype(("S", max_count))) 2071 else: -> 2072 return self.to_RegularArray()._to_backend_array(allow_missing, backend) File [~/coffea-gpu/awkward/src/awkward/contents/listoffsetarray.py:283](https://analytics-hub.fnal.gov/user/lagray/lab/tree/coffea-gpu/coffea-gpu/awkward/src/awkward/contents/listoffsetarray.py#line=282), in ListOffsetArray.to_RegularArray(self) 278 _size = Index64.empty(1, self._backend.index_nplike) 279 assert ( 280 _size.nplike is self._backend.index_nplike 281 and self._offsets.nplike is self._backend.index_nplike 282 ) --> 283 self._backend.maybe_kernel_error( 284 self._backend[ 285 "awkward_ListOffsetArray_toRegularArray", 286 _size.dtype.type, 287 self._offsets.dtype.type, 288 ]( 289 _size.data, 290 self._offsets.data, 291 self._offsets.length, 292 ) 293 ) 294 size = self._backend.index_nplike.index_as_shape_item(_size[0]) 295 length = self._offsets.length - 1 File [~/coffea-gpu/awkward/src/awkward/_backends/backend.py:67](https://analytics-hub.fnal.gov/user/lagray/lab/tree/coffea-gpu/coffea-gpu/awkward/src/awkward/_backends/backend.py#line=66), in Backend.maybe_kernel_error(self, error) 65 return 66 else: ---> 67 raise ValueError(self.format_kernel_error(error)) ValueError: cannot convert to RegularArray because subarray lengths are not regular (in compiled code: https://github.com/scikit-hep/awkward/blob/awkward-cpp-35/awkward-cpp/src/cpu-kernels/awkward_ListOffsetArray_toRegularArray.cpp#L22) This error occurred while calling numpy.asarray( <Array [[0.2120250193289826, ...], ...] type='50 * var * float64'> dtype = None ) 

Metadata

Metadata

Assignees

Labels

bug (unverified)The problem described would be a bug, but needs to be triaged

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions