From e4c6c9ccfb04d3aa97add0549e866c945a3aa08b Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Thu, 1 Aug 2024 14:36:34 +0200 Subject: [PATCH] add test --- tests-cuda/test_3191_numba_cuda_vectorize.py | 35 ++++++++++++++++++++ 1 file changed, 35 insertions(+) create mode 100644 tests-cuda/test_3191_numba_cuda_vectorize.py diff --git a/tests-cuda/test_3191_numba_cuda_vectorize.py b/tests-cuda/test_3191_numba_cuda_vectorize.py new file mode 100644 index 0000000000..7aaf5b9141 --- /dev/null +++ b/tests-cuda/test_3191_numba_cuda_vectorize.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +import cupy as cp +import numba as nb + +import awkward as ak + +ak.numba.register_and_check() + + +@nb.vectorize() +def _square(x): + return x * x + + +@nb.vectorize( + 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) + + +def test_square(): + values2_cpu = _square(ak.to_backend(values, "cpu")) + print(values2_cpu) + + +def test_square_cuda(): + values2 = _square_cuda(values) + print(values2)