Skip to content

Commit 772fd97

Browse files
committed
Add cupy tests
1 parent 354a4eb commit 772fd97

File tree

1 file changed

+23
-0
lines changed

1 file changed

+23
-0
lines changed

tests/test_cupy.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
import unittest
2+
3+
import cupy as cp
4+
5+
from common import gpu_test
6+
7+
8+
class TestCupy(unittest.TestCase):
9+
@gpu_test
10+
def test_kernel(self):
11+
x = cp.arange(6, dtype='f').reshape(2, 3)
12+
y = cp.arange(3, dtype='f')
13+
kernel = cp.ElementwiseKernel(
14+
'float32 x, float32 y', 'float32 z',
15+
'''if (x - 2 > y) {
16+
z = x * y;
17+
} else {
18+
z = x + y;
19+
}''',
20+
'my_kernel')
21+
r = kernel(x, y)
22+
23+
self.assertEqual((2, 3), r.shape)

0 commit comments

Comments
 (0)