warp-lang 0.10.1__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/tests/test_array.py

@@ -5,14 +5,12 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-# include parent path
+import unittest
+
 import numpy as np
-import math
 
 import warp as wp
-from warp.tests.test_base import *
-
-import unittest
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -397,7 +395,7 @@ def test_slicing(test, device):
     assert_array_equal(wp_arr[:5], wp.array(np_arr[:5], dtype=int, device=device))
     assert_array_equal(wp_arr[1:5], wp.array(np_arr[1:5], dtype=int, device=device))
     assert_array_equal(wp_arr[-9:-5:1], wp.array(np_arr[-9:-5:1], dtype=int, device=device))
-    assert_array_equal(wp_arr[:5,], wp.array(np_arr[:5], dtype=int, device=device))
+    assert_array_equal(wp_arr[:5,], wp.array(np_arr[:5], dtype=int, device=device))  # noqa: E231
 
 
 def test_view(test, device):
@@ -738,7 +736,10 @@ def test_fill_matrix(test, device):
             assert_np_equal(a4.numpy(), np.zeros((*a4.shape, *mat_shape), dtype=nptype))
 
             # matrix values can be passed as a 1d numpy array, 2d numpy array, flat list, nested list, or Warp matrix instance
-            fill_arr1 = np.arange(mat_len, dtype=nptype)
+            if wptype != wp.bool:
+                fill_arr1 = np.arange(mat_len, dtype=nptype)
+            else:
+                fill_arr1 = np.ones(mat_len, dtype=nptype)
             fill_arr2 = fill_arr1.reshape(mat_shape)
             fill_list1 = list(fill_arr1)
             fill_list2 = [list(row) for row in fill_arr2]
@@ -1295,7 +1296,10 @@ def test_full_matrix(test, device):
                 assert_np_equal(na, np.full(a.size * mattype._length_, 42, dtype=nptype).reshape(npshape))
 
                 # fill with 1d numpy array and specific dtype
-                fill_arr1d = np.arange(mattype._length_, dtype=nptype)
+                if wptype != wp.bool:
+                    fill_arr1d = np.arange(mattype._length_, dtype=nptype)
+                else:
+                    fill_arr1d = np.ones(mattype._length_, dtype=nptype)
                 a = wp.full(shape, fill_arr1d, dtype=mattype, device=device)
                 na = a.numpy()
 
@@ -1448,16 +1452,17 @@ def test_full_struct(test, device):
 
 
 def test_round_trip(test, device):
+    rng = np.random.default_rng(123)
     dim_x = 4
 
     for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
-        a_np = np.random.randn(dim_x).astype(nptype)
+        a_np = rng.standard_normal(size=dim_x).astype(nptype)
         a = wp.array(a_np, device=device)
         test.assertEqual(a.dtype, wptype)
 
         assert_np_equal(a.numpy(), a_np)
 
-        v_np = np.random.randn(dim_x, 3).astype(nptype)
+        v_np = rng.standard_normal(size=(dim_x, 3)).astype(nptype)
         v = wp.array(v_np, dtype=wp.types.vector(3, wptype), device=device)
 
         assert_np_equal(v.numpy(), v_np)
@@ -1695,6 +1700,7 @@ def test_to_list_struct(test, device):
         a1: wp.array(dtype=int)
         a2: wp.array2d(dtype=float)
         a3: wp.array3d(dtype=wp.float16)
+        bool: wp.bool
 
     dim = 3
 
@@ -1714,6 +1720,7 @@ def test_to_list_struct(test, device):
     s.a1 = wp.empty(1, dtype=int, device=device)
     s.a2 = wp.empty((1, 1), dtype=float, device=device)
     s.a3 = wp.empty((1, 1, 1), dtype=wp.float16, device=device)
+    s.bool = True
 
     for ndim in range(1, 5):
         shape = (dim,) * ndim
@@ -1731,6 +1738,7 @@ def test_to_list_struct(test, device):
             test.assertEqual(l[i].mi, s.mi)
             test.assertEqual(l[i].mf, s.mf)
             test.assertEqual(l[i].mh, s.mh)
+            test.assertEqual(l[i].bool, s.bool)
             test.assertEqual(l[i].inner.h, s.inner.h)
             test.assertEqual(l[i].inner.v, s.inner.v)
             test.assertEqual(l[i].a1.dtype, s.a1.dtype)
@@ -1741,46 +1749,6 @@ def test_to_list_struct(test, device):
             test.assertEqual(l[i].a3.ndim, s.a3.ndim)
 
 
-def test_large_arrays_slow(test, device):
-    # The goal of this test is to use arrays just large enough to know
-    # if there's a flaw in handling arrays with more than 2**31-1 elements
-    # Unfortunately, it takes a long time to run so it won't be run automatically
-    # without changes to support how frequently a test may be run
-    total_elements = 2**31 + 8
-
-    # 1-D to 4-D arrays: test zero_, fill_, then zero_ for scalar data types:
-    for total_dims in range(1, 5):
-        dim_x = math.ceil(total_elements ** (1 / total_dims))
-        shape_tuple = tuple([dim_x] * total_dims)
-
-        for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
-            a1 = wp.zeros(shape_tuple, dtype=wptype, device=device)
-            assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
-
-            a1.fill_(127)
-            assert_np_equal(a1.numpy(), 127 * np.ones_like(a1.numpy()))
-
-            a1.zero_()
-            assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
-
-
-def test_large_arrays_fast(test, device):
-    # A truncated version of test_large_arrays_slow meant to catch basic errors
-    total_elements = 2**31 + 8
-
-    nptype = np.dtype(np.int8)
-    wptype = wp.types.np_dtype_to_warp_type[nptype]
-
-    a1 = wp.zeros((total_elements,), dtype=wptype, device=device)
-    assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
-
-    a1.fill_(127)
-    assert_np_equal(a1.numpy(), 127 * np.ones_like(a1.numpy()))
-
-    a1.zero_()
-    assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
-
-
 @wp.kernel
 def kernel_array_to_bool(array_null: wp.array(dtype=float), array_valid: wp.array(dtype=float)):
     if not array_null:
@@ -1969,54 +1937,201 @@ def test_array_of_structs_roundtrip(test, device):
     assert_np_equal(a.numpy(), expected)
 
 
-def register(parent):
-    devices = get_test_devices()
-
-    class TestArray(parent):
-        pass
-
-    add_function_test(TestArray, "test_shape", test_shape, devices=devices)
-    add_function_test(TestArray, "test_flatten", test_flatten, devices=devices)
-    add_function_test(TestArray, "test_reshape", test_reshape, devices=devices)
-    add_function_test(TestArray, "test_slicing", test_slicing, devices=devices)
-    add_function_test(TestArray, "test_transpose", test_transpose, devices=devices)
-    add_function_test(TestArray, "test_view", test_view, devices=devices)
-
-    add_function_test(TestArray, "test_1d_array", test_1d, devices=devices)
-    add_function_test(TestArray, "test_2d_array", test_2d, devices=devices)
-    add_function_test(TestArray, "test_3d_array", test_3d, devices=devices)
-    add_function_test(TestArray, "test_4d_array", test_4d, devices=devices)
-    add_function_test(TestArray, "test_4d_array_transposed", test_4d_transposed, devices=devices)
-
-    add_function_test(TestArray, "test_fill_scalar", test_fill_scalar, devices=devices)
-    add_function_test(TestArray, "test_fill_vector", test_fill_vector, devices=devices)
-    add_function_test(TestArray, "test_fill_matrix", test_fill_matrix, devices=devices)
-    add_function_test(TestArray, "test_fill_struct", test_fill_struct, devices=devices)
-    add_function_test(TestArray, "test_fill_slices", test_fill_slices, devices=devices)
-    add_function_test(TestArray, "test_full_scalar", test_full_scalar, devices=devices)
-    add_function_test(TestArray, "test_full_vector", test_full_vector, devices=devices)
-    add_function_test(TestArray, "test_full_matrix", test_full_matrix, devices=devices)
-    add_function_test(TestArray, "test_full_struct", test_full_struct, devices=devices)
-    add_function_test(TestArray, "test_empty_array", test_empty_array, devices=devices)
-    add_function_test(TestArray, "test_empty_from_numpy", test_empty_from_numpy, devices=devices)
-    add_function_test(TestArray, "test_empty_from_list", test_empty_from_list, devices=devices)
-    add_function_test(TestArray, "test_to_list_scalar", test_to_list_scalar, devices=devices)
-    add_function_test(TestArray, "test_to_list_vector", test_to_list_vector, devices=devices)
-    add_function_test(TestArray, "test_to_list_matrix", test_to_list_matrix, devices=devices)
-    add_function_test(TestArray, "test_to_list_struct", test_to_list_struct, devices=devices)
-
-    add_function_test(TestArray, "test_lower_bound", test_lower_bound, devices=devices)
-    add_function_test(TestArray, "test_round_trip", test_round_trip, devices=devices)
-    add_function_test(TestArray, "test_large_arrays_fast", test_large_arrays_fast, devices=devices)
-    add_function_test(TestArray, "test_array_to_bool", test_array_to_bool, devices=devices)
-    add_function_test(TestArray, "test_array_of_structs", test_array_of_structs, devices=devices)
-    add_function_test(TestArray, "test_array_of_structs_grad", test_array_of_structs_grad, devices=devices)
-    add_function_test(TestArray, "test_array_of_structs_from_numpy", test_array_of_structs_from_numpy, devices=devices)
-    add_function_test(TestArray, "test_array_of_structs_roundtrip", test_array_of_structs_roundtrip, devices=devices)
-
-    return TestArray
+def test_array_from_numpy(test, device):
+    arr = np.array((1.0, 2.0, 3.0), dtype=float)
+
+    result = wp.from_numpy(arr)
+    expected = wp.array((1.0, 2.0, 3.0), dtype=wp.float32, shape=(3,))
+    assert_np_equal(result.numpy(), expected.numpy())
+
+    result = wp.from_numpy(arr, dtype=wp.vec3)
+    expected = wp.array(((1.0, 2.0, 3.0),), dtype=wp.vec3, shape=(1,))
+    assert_np_equal(result.numpy(), expected.numpy())
+
+    # --------------------------------------------------------------------------
+
+    arr = np.array(((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)), dtype=float)
+
+    result = wp.from_numpy(arr)
+    expected = wp.array(((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)), dtype=wp.vec3, shape=(2,))
+    assert_np_equal(result.numpy(), expected.numpy())
+
+    result = wp.from_numpy(arr, dtype=wp.float32)
+    expected = wp.array(((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)), dtype=wp.float32, shape=(2, 3))
+    assert_np_equal(result.numpy(), expected.numpy())
+
+    result = wp.from_numpy(arr, dtype=wp.float32, shape=(6,))
+    expected = wp.array((1.0, 2.0, 3.0, 4.0, 5.0, 6.0), dtype=wp.float32, shape=(6,))
+    assert_np_equal(result.numpy(), expected.numpy())
+
+    # --------------------------------------------------------------------------
+
+    arr = np.array(
+        (
+            (
+                (1.0, 2.0, 3.0, 4.0),
+                (2.0, 3.0, 4.0, 5.0),
+                (3.0, 4.0, 5.0, 6.0),
+                (4.0, 5.0, 6.0, 7.0),
+            ),
+            (
+                (2.0, 3.0, 4.0, 5.0),
+                (3.0, 4.0, 5.0, 6.0),
+                (4.0, 5.0, 6.0, 7.0),
+                (5.0, 6.0, 7.0, 8.0),
+            ),
+        ),
+        dtype=float,
+    )
+
+    result = wp.from_numpy(arr)
+    expected = wp.array(
+        (
+            (
+                (1.0, 2.0, 3.0, 4.0),
+                (2.0, 3.0, 4.0, 5.0),
+                (3.0, 4.0, 5.0, 6.0),
+                (4.0, 5.0, 6.0, 7.0),
+            ),
+            (
+                (2.0, 3.0, 4.0, 5.0),
+                (3.0, 4.0, 5.0, 6.0),
+                (4.0, 5.0, 6.0, 7.0),
+                (5.0, 6.0, 7.0, 8.0),
+            ),
+        ),
+        dtype=wp.mat44,
+        shape=(2,),
+    )
+    assert_np_equal(result.numpy(), expected.numpy())
+
+    result = wp.from_numpy(arr, dtype=wp.float32)
+    expected = wp.array(
+        (
+            (
+                (1.0, 2.0, 3.0, 4.0),
+                (2.0, 3.0, 4.0, 5.0),
+                (3.0, 4.0, 5.0, 6.0),
+                (4.0, 5.0, 6.0, 7.0),
+            ),
+            (
+                (2.0, 3.0, 4.0, 5.0),
+                (3.0, 4.0, 5.0, 6.0),
+                (4.0, 5.0, 6.0, 7.0),
+                (5.0, 6.0, 7.0, 8.0),
+            ),
+        ),
+        dtype=wp.float32,
+        shape=(2, 4, 4),
+    )
+    assert_np_equal(result.numpy(), expected.numpy())
+
+    result = wp.from_numpy(arr, dtype=wp.vec4)
+    expected = wp.array(
+        (
+            (1.0, 2.0, 3.0, 4.0),
+            (2.0, 3.0, 4.0, 5.0),
+            (3.0, 4.0, 5.0, 6.0),
+            (4.0, 5.0, 6.0, 7.0),
+            (2.0, 3.0, 4.0, 5.0),
+            (3.0, 4.0, 5.0, 6.0),
+            (4.0, 5.0, 6.0, 7.0),
+            (5.0, 6.0, 7.0, 8.0),
+        ),
+        dtype=wp.vec4,
+        shape=(8,),
+    )
+    assert_np_equal(result.numpy(), expected.numpy())
+
+    result = wp.from_numpy(arr, dtype=wp.float32, shape=(32,))
+    expected = wp.array(
+        (
+            1.0,
+            2.0,
+            3.0,
+            4.0,
+            2.0,
+            3.0,
+            4.0,
+            5.0,
+            3.0,
+            4.0,
+            5.0,
+            6.0,
+            4.0,
+            5.0,
+            6.0,
+            7.0,
+            2.0,
+            3.0,
+            4.0,
+            5.0,
+            3.0,
+            4.0,
+            5.0,
+            6.0,
+            4.0,
+            5.0,
+            6.0,
+            7.0,
+            5.0,
+            6.0,
+            7.0,
+            8.0,
+        ),
+        dtype=wp.float32,
+        shape=(32,),
+    )
+    assert_np_equal(result.numpy(), expected.numpy())
+
+
+devices = get_test_devices()
+
+
+class TestArray(unittest.TestCase):
+    pass
+
+
+add_function_test(TestArray, "test_shape", test_shape, devices=devices)
+add_function_test(TestArray, "test_flatten", test_flatten, devices=devices)
+add_function_test(TestArray, "test_reshape", test_reshape, devices=devices)
+add_function_test(TestArray, "test_slicing", test_slicing, devices=devices)
+add_function_test(TestArray, "test_transpose", test_transpose, devices=devices)
+add_function_test(TestArray, "test_view", test_view, devices=devices)
+
+add_function_test(TestArray, "test_1d_array", test_1d, devices=devices)
+add_function_test(TestArray, "test_2d_array", test_2d, devices=devices)
+add_function_test(TestArray, "test_3d_array", test_3d, devices=devices)
+add_function_test(TestArray, "test_4d_array", test_4d, devices=devices)
+add_function_test(TestArray, "test_4d_array_transposed", test_4d_transposed, devices=devices)
+
+add_function_test(TestArray, "test_fill_scalar", test_fill_scalar, devices=devices)
+add_function_test(TestArray, "test_fill_vector", test_fill_vector, devices=devices)
+add_function_test(TestArray, "test_fill_matrix", test_fill_matrix, devices=devices)
+add_function_test(TestArray, "test_fill_struct", test_fill_struct, devices=devices)
+add_function_test(TestArray, "test_fill_slices", test_fill_slices, devices=devices)
+add_function_test(TestArray, "test_full_scalar", test_full_scalar, devices=devices)
+add_function_test(TestArray, "test_full_vector", test_full_vector, devices=devices)
+add_function_test(TestArray, "test_full_matrix", test_full_matrix, devices=devices)
+add_function_test(TestArray, "test_full_struct", test_full_struct, devices=devices)
+add_function_test(TestArray, "test_empty_array", test_empty_array, devices=devices)
+add_function_test(TestArray, "test_empty_from_numpy", test_empty_from_numpy, devices=devices)
+add_function_test(TestArray, "test_empty_from_list", test_empty_from_list, devices=devices)
+add_function_test(TestArray, "test_to_list_scalar", test_to_list_scalar, devices=devices)
+add_function_test(TestArray, "test_to_list_vector", test_to_list_vector, devices=devices)
+add_function_test(TestArray, "test_to_list_matrix", test_to_list_matrix, devices=devices)
+add_function_test(TestArray, "test_to_list_struct", test_to_list_struct, devices=devices)
+
+add_function_test(TestArray, "test_lower_bound", test_lower_bound, devices=devices)
+add_function_test(TestArray, "test_round_trip", test_round_trip, devices=devices)
+add_function_test(TestArray, "test_array_to_bool", test_array_to_bool, devices=devices)
+add_function_test(TestArray, "test_array_of_structs", test_array_of_structs, devices=devices)
+add_function_test(TestArray, "test_array_of_structs_grad", test_array_of_structs_grad, devices=devices)
+add_function_test(TestArray, "test_array_of_structs_from_numpy", test_array_of_structs_from_numpy, devices=devices)
+add_function_test(TestArray, "test_array_of_structs_roundtrip", test_array_of_structs_roundtrip, devices=devices)
+add_function_test(TestArray, "test_array_from_numpy", test_array_from_numpy, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_array_reduce.py

@@ -1,8 +1,17 @@
+# Copyright (c) 2023 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import unittest
+
 import numpy as np
-import warp as wp
 
-from warp.utils import array_sum, array_inner
-from warp.tests.test_base import *
+import warp as wp
+from warp.tests.unittest_utils import *
+from warp.utils import array_inner, array_sum
 
 wp.init()
 
@@ -11,9 +20,11 @@ def make_test_array_sum(dtype):
     N = 1000
 
     def test_array_sum(test, device):
+        rng = np.random.default_rng(123)
+
         cols = wp.types.type_length(dtype)
 
-        values_np = np.random.rand(N, cols)
+        values_np = rng.random(size=(N, cols))
         values = wp.array(values_np, device=device, dtype=dtype)
 
         vsum = array_sum(values)
@@ -32,7 +43,9 @@ def make_test_array_sum_axis(dtype):
     N = I * J * K
 
     def test_array_sum(test, device):
-        values_np = np.random.rand(I, J, K)
+        rng = np.random.default_rng(123)
+
+        values_np = rng.random(size=(I, J, K))
         values = wp.array(values_np, shape=(I, J, K), device=device, dtype=dtype)
 
         for axis in range(3):
@@ -44,14 +57,24 @@ def make_test_array_sum_axis(dtype):
     return test_array_sum
 
 
+def test_array_sum_empty(test, device):
+    values = wp.array([], device=device, dtype=wp.vec2)
+    assert_np_equal(array_sum(values), np.zeros(2))
+
+    values = wp.array([], shape=(0, 3), device=device, dtype=float)
+    assert_np_equal(array_sum(values, axis=0).numpy(), np.zeros(3))
+
+
 def make_test_array_inner(dtype):
     N = 1000
 
     def test_array_inner(test, device):
+        rng = np.random.default_rng(123)
+
         cols = wp.types.type_length(dtype)
 
-        a_np = np.random.rand(N, cols)
-        b_np = np.random.rand(N, cols)
+        a_np = rng.random(size=(N, cols))
+        b_np = rng.random(size=(N, cols))
 
         a = wp.array(a_np, device=device, dtype=dtype)
         b = wp.array(b_np, device=device, dtype=dtype)
@@ -72,8 +95,10 @@ def make_test_array_inner_axis(dtype):
     N = I * J * K
 
     def test_array_inner(test, device):
-        a_np = np.random.rand(I, J, K)
-        b_np = np.random.rand(I, J, K)
+        rng = np.random.default_rng(123)
+
+        a_np = rng.random(size=(I, J, K))
+        b_np = rng.random(size=(I, J, K))
 
         a = wp.array(a_np, shape=(I, J, K), device=device, dtype=dtype)
         b = wp.array(b_np, shape=(I, J, K), device=device, dtype=dtype)
@@ -93,24 +118,33 @@ def make_test_array_inner_axis(dtype):
     return test_array_inner
 
 
-def register(parent):
-    devices = get_test_devices()
+def test_array_inner_empty(test, device):
+    values = wp.array([], device=device, dtype=wp.vec2)
+    test.assertEqual(array_inner(values, values), 0.0)
+
+    values = wp.array([], shape=(0, 3), device=device, dtype=float)
+    assert_np_equal(array_inner(values, values, axis=0).numpy(), np.zeros(3))
+
+
+devices = get_test_devices()
+
 
-    class TestArraySym(parent):
-        pass
+class TestArrayReduce(unittest.TestCase):
+    pass
 
-    add_function_test(TestArraySym, "test_array_sum_double", make_test_array_sum(wp.float64), devices=devices)
-    add_function_test(TestArraySym, "test_array_sum_vec3", make_test_array_sum(wp.vec3), devices=devices)
-    add_function_test(TestArraySym, "test_array_sum_axis_float", make_test_array_sum_axis(wp.float32), devices=devices)
-    add_function_test(TestArraySym, "test_array_inner_double", make_test_array_inner(wp.float64), devices=devices)
-    add_function_test(TestArraySym, "test_array_inner_vec3", make_test_array_inner(wp.vec3), devices=devices)
-    add_function_test(
-        TestArraySym, "test_array_inner_axis_float", make_test_array_inner_axis(wp.float32), devices=devices
-    )
 
-    return TestArraySym
+add_function_test(TestArrayReduce, "test_array_sum_double", make_test_array_sum(wp.float64), devices=devices)
+add_function_test(TestArrayReduce, "test_array_sum_vec3", make_test_array_sum(wp.vec3), devices=devices)
+add_function_test(TestArrayReduce, "test_array_sum_axis_float", make_test_array_sum_axis(wp.float32), devices=devices)
+add_function_test(TestArrayReduce, "test_array_sum_empty", test_array_sum_empty, devices=devices)
+add_function_test(TestArrayReduce, "test_array_inner_double", make_test_array_inner(wp.float64), devices=devices)
+add_function_test(TestArrayReduce, "test_array_inner_vec3", make_test_array_inner(wp.vec3), devices=devices)
+add_function_test(
+    TestArrayReduce, "test_array_inner_axis_float", make_test_array_inner_axis(wp.float32), devices=devices
+)
+add_function_test(TestArrayReduce, "test_array_inner_empty", test_array_inner_empty, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_atomic.py

@@ -5,14 +5,12 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-# include parent path
+import unittest
+
 import numpy as np
-import math
 
 import warp as wp
-from warp.tests.test_base import *
-
-import unittest
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -34,8 +32,7 @@ def make_atomic_test(type):
     # register a custom kernel (no decorator) function
     # this lets us register the same function definition
     # against multiple symbols, with different arg types
-    module = wp.get_module(test_atomic_kernel.__module__)
-    kernel = wp.Kernel(func=test_atomic_kernel, key=f"test_atomic_{type.__name__}_kernel", module=module)
+    kernel = wp.Kernel(func=test_atomic_kernel, key=f"test_atomic_{type.__name__}_kernel")
 
     def test_atomic(test, device):
         n = 1024
@@ -54,20 +51,60 @@ def make_atomic_test(type):
             base = rng.random(size=(1, *type._shape_), dtype=float)
             val = rng.random(size=(n, *type._shape_), dtype=float)
 
-        add_array = wp.array(base, dtype=type, device=device)
-        min_array = wp.array(base, dtype=type, device=device)
-        max_array = wp.array(base, dtype=type, device=device)
+        add_array = wp.array(base, dtype=type, device=device, requires_grad=True)
+        min_array = wp.array(base, dtype=type, device=device, requires_grad=True)
+        max_array = wp.array(base, dtype=type, device=device, requires_grad=True)
+        add_array.zero_()
+        min_array.fill_(10000)
+        max_array.fill_(-10000)
 
-        val_array = wp.array(val, dtype=type, device=device)
+        val_array = wp.array(val, dtype=type, device=device, requires_grad=True)
 
-        wp.launch(kernel, n, inputs=[add_array, min_array, max_array, val_array], device=device)
-
-        val = np.append(val, [base[0]], axis=0)
+        tape = wp.Tape()
+        with tape:
+            wp.launch(kernel, n, inputs=[add_array, min_array, max_array, val_array], device=device)
 
         assert_np_equal(add_array.numpy(), np.sum(val, axis=0), tol=1.0e-2)
         assert_np_equal(min_array.numpy(), np.min(val, axis=0), tol=1.0e-2)
         assert_np_equal(max_array.numpy(), np.max(val, axis=0), tol=1.0e-2)
 
+        if type != wp.int32:
+            add_array.grad.fill_(1)
+            tape.backward()
+            assert_np_equal(val_array.grad.numpy(), np.ones_like(val))
+            tape.zero()
+
+            min_array.grad.fill_(1)
+            tape.backward()
+            min_grad_array = np.zeros_like(val)
+            argmin = val.argmin(axis=0)
+            if val.ndim == 1:
+                min_grad_array[argmin] = 1
+            elif val.ndim == 2:
+                for i in range(val.shape[1]):
+                    min_grad_array[argmin[i], i] = 1
+            elif val.ndim == 3:
+                for i in range(val.shape[1]):
+                    for j in range(val.shape[2]):
+                        min_grad_array[argmin[i, j], i, j] = 1
+            assert_np_equal(val_array.grad.numpy(), min_grad_array)
+            tape.zero()
+
+            max_array.grad.fill_(1)
+            tape.backward()
+            max_grad_array = np.zeros_like(val)
+            argmax = val.argmax(axis=0)
+            if val.ndim == 1:
+                max_grad_array[argmax] = 1
+            elif val.ndim == 2:
+                for i in range(val.shape[1]):
+                    max_grad_array[argmax[i], i] = 1
+            elif val.ndim == 3:
+                for i in range(val.shape[1]):
+                    for j in range(val.shape[2]):
+                        max_grad_array[argmax[i, j], i, j] = 1
+            assert_np_equal(val_array.grad.numpy(), max_grad_array)
+
     return test_atomic
 
 
@@ -82,24 +119,23 @@ test_atomic_mat33 = make_atomic_test(wp.mat33)
 test_atomic_mat44 = make_atomic_test(wp.mat44)
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestAtomic(parent):
-        pass
+class TestAtomic(unittest.TestCase):
+    pass
 
-    add_function_test(TestAtomic, "test_atomic_int", test_atomic_int, devices=devices)
-    add_function_test(TestAtomic, "test_atomic_float", test_atomic_float, devices=devices)
-    add_function_test(TestAtomic, "test_atomic_vec2", test_atomic_vec2, devices=devices)
-    add_function_test(TestAtomic, "test_atomic_vec3", test_atomic_vec3, devices=devices)
-    add_function_test(TestAtomic, "test_atomic_vec4", test_atomic_vec4, devices=devices)
-    add_function_test(TestAtomic, "test_atomic_mat22", test_atomic_mat22, devices=devices)
-    add_function_test(TestAtomic, "test_atomic_mat33", test_atomic_mat33, devices=devices)
-    add_function_test(TestAtomic, "test_atomic_mat44", test_atomic_mat44, devices=devices)
 
-    return TestAtomic
+add_function_test(TestAtomic, "test_atomic_int", test_atomic_int, devices=devices)
+add_function_test(TestAtomic, "test_atomic_float", test_atomic_float, devices=devices)
+add_function_test(TestAtomic, "test_atomic_vec2", test_atomic_vec2, devices=devices)
+add_function_test(TestAtomic, "test_atomic_vec3", test_atomic_vec3, devices=devices)
+add_function_test(TestAtomic, "test_atomic_vec4", test_atomic_vec4, devices=devices)
+add_function_test(TestAtomic, "test_atomic_mat22", test_atomic_mat22, devices=devices)
+add_function_test(TestAtomic, "test_atomic_mat33", test_atomic_mat33, devices=devices)
+add_function_test(TestAtomic, "test_atomic_mat44", test_atomic_mat44, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)