Fix Numba serialization when strides is None
@@ -18,9 +18,18 @@ def serialize_numba_ndarray(x):
@cuda_deserialize.register(numba.cuda.devicearray.DeviceNDArray)
def deserialize_numba_ndarray(header, frames):
(frame,) = frames
+ shape = header["shape"]
+ strides = header["strides"]
+
+ # Starting with __cuda_array_interface__ version 2, strides can be None,
+ # meaning the array is C-contiguous, so we have to calculate it.
+ if strides is None:
+ itemsize = np.dtype(header["typestr"]).itemsize
+ strides = tuple((np.cumprod((1,) + shape[:0:-1]) * itemsize).tolist())
+
arr = numba.cuda.devicearray.DeviceNDArray(
- header["shape"],
- header["strides"],
+ shape,
+ strides,
np.dtype(header["typestr"]),
gpu_data=numba.cuda.as_cuda_array(frame).gpu_data,
)
