warp-lang 1.0.0b5__py3-none-manylinux2014_x86_64.whl → 1.0.0b6__py3-none-manylinux2014_x86_64.whl

warp/tests/test_grad.py

@@ -6,11 +6,12 @@
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
 import unittest
+from typing import Any
 
 import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -66,26 +67,26 @@ def test_for_loop_grad(test, device):
 
 
 def test_for_loop_graph_grad(test, device):
+    wp.load_module(device=device)
+
     n = 32
     val = np.ones(n, dtype=np.float32)
 
     x = wp.array(val, device=device, requires_grad=True)
     sum = wp.zeros(1, dtype=wp.float32, device=device, requires_grad=True)
 
-    wp.force_load()
-
-    wp.capture_begin()
-
-    tape = wp.Tape()
-    with tape:
-        wp.launch(for_loop_grad, dim=1, inputs=[n, x, sum], device=device)
-
-    tape.backward(loss=sum)
+    wp.capture_begin(device, force_module_load=False)
+    try:
+        tape = wp.Tape()
+        with tape:
+            wp.launch(for_loop_grad, dim=1, inputs=[n, x, sum], device=device)
 
-    graph = wp.capture_end()
+        tape.backward(loss=sum)
+    finally:
+        graph = wp.capture_end(device)
 
     wp.capture_launch(graph)
-    wp.synchronize()
+    wp.synchronize_device(device)
 
     # ensure forward pass outputs persist
     assert_np_equal(sum.numpy(), 2.0 * np.sum(x.numpy()))
@@ -93,7 +94,7 @@ def test_for_loop_graph_grad(test, device):
     assert_np_equal(x.grad.numpy(), 2.0 * val)
 
     wp.capture_launch(graph)
-    wp.synchronize()
+    wp.synchronize_device(device)
 
 
 @wp.kernel
@@ -560,30 +561,80 @@ def test_name_clash(test, device):
         assert_np_equal(input_b.grad.numpy(), np.array([1.0, 1.0, 0.0]))
 
 
-def register(parent):
-    devices = get_test_devices()
+@wp.struct
+class NestedStruct:
+    v: wp.vec2
+
+
+@wp.struct
+class ParentStruct:
+    a: float
+    n: NestedStruct
+
+
+@wp.func
+def noop(a: Any):
+    pass
+
+
+@wp.func
+def sum2(v: wp.vec2):
+    return v[0] + v[1]
+
+
+@wp.kernel
+def test_struct_attribute_gradient_kernel(src: wp.array(dtype=float), res: wp.array(dtype=float)):
+    tid = wp.tid()
+
+    p = ParentStruct(src[tid], NestedStruct(wp.vec2(2.0 * src[tid])))
+
+    # test that we are not losing gradients when accessing attributes
+    noop(p.a)
+    noop(p.n)
+    noop(p.n.v)
+
+    res[tid] = p.a + sum2(p.n.v)
+
+
+def test_struct_attribute_gradient(test_case, device):
+    src = wp.array([1], dtype=float, requires_grad=True)
+    res = wp.empty_like(src)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(test_struct_attribute_gradient_kernel, dim=1, inputs=[src, res])
+
+    res.grad.fill_(1.0)
+    tape.backward()
+
+    test_case.assertEqual(src.grad.numpy()[0], 5.0)
+
+
+devices = get_test_devices()
+
 
-    class TestGrad(parent):
-        pass
+class TestGrad(unittest.TestCase):
+    pass
 
-    # add_function_test(TestGrad, "test_while_loop_grad", test_while_loop_grad, devices=devices)
-    add_function_test(TestGrad, "test_for_loop_nested_for_grad", test_for_loop_nested_for_grad, devices=devices)
-    add_function_test(TestGrad, "test_scalar_grad", test_scalar_grad, devices=devices)
-    add_function_test(TestGrad, "test_for_loop_grad", test_for_loop_grad, devices=devices)
-    add_function_test(TestGrad, "test_for_loop_graph_grad", test_for_loop_graph_grad, devices=wp.get_cuda_devices())
-    add_function_test(TestGrad, "test_for_loop_nested_if_grad", test_for_loop_nested_if_grad, devices=devices)
-    add_function_test(TestGrad, "test_preserve_outputs_grad", test_preserve_outputs_grad, devices=devices)
-    add_function_test(TestGrad, "test_vector_math_grad", test_vector_math_grad, devices=devices)
-    add_function_test(TestGrad, "test_matrix_math_grad", test_matrix_math_grad, devices=devices)
-    add_function_test(TestGrad, "test_3d_math_grad", test_3d_math_grad, devices=devices)
-    add_function_test(TestGrad, "test_multi_valued_function_grad", test_multi_valued_function_grad, devices=devices)
-    add_function_test(TestGrad, "test_mesh_grad", test_mesh_grad, devices=devices)
-    add_function_test(TestGrad, "test_name_clash", test_name_clash, devices=devices)
 
-    return TestGrad
+# add_function_test(TestGrad, "test_while_loop_grad", test_while_loop_grad, devices=devices)
+add_function_test(TestGrad, "test_for_loop_nested_for_grad", test_for_loop_nested_for_grad, devices=devices)
+add_function_test(TestGrad, "test_scalar_grad", test_scalar_grad, devices=devices)
+add_function_test(TestGrad, "test_for_loop_grad", test_for_loop_grad, devices=devices)
+add_function_test(
+    TestGrad, "test_for_loop_graph_grad", test_for_loop_graph_grad, devices=get_unique_cuda_test_devices()
+)
+add_function_test(TestGrad, "test_for_loop_nested_if_grad", test_for_loop_nested_if_grad, devices=devices)
+add_function_test(TestGrad, "test_preserve_outputs_grad", test_preserve_outputs_grad, devices=devices)
+add_function_test(TestGrad, "test_vector_math_grad", test_vector_math_grad, devices=devices)
+add_function_test(TestGrad, "test_matrix_math_grad", test_matrix_math_grad, devices=devices)
+add_function_test(TestGrad, "test_3d_math_grad", test_3d_math_grad, devices=devices)
+add_function_test(TestGrad, "test_multi_valued_function_grad", test_multi_valued_function_grad, devices=devices)
+add_function_test(TestGrad, "test_mesh_grad", test_mesh_grad, devices=devices)
+add_function_test(TestGrad, "test_name_clash", test_name_clash, devices=devices)
+add_function_test(TestGrad, "test_struct_attribute_gradient", test_struct_attribute_gradient, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2, failfast=False)

warp/tests/test_grad_customs.py

@@ -10,7 +10,7 @@ import unittest
 import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -160,19 +160,17 @@ def test_custom_overload_grad(test, device):
     # fmt: on
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
 
-    class TestGradCustoms(parent):
-        pass
 
-    add_function_test(TestGradCustoms, "test_custom_replay_grad", test_custom_replay_grad, devices=devices)
-    add_function_test(TestGradCustoms, "test_custom_overload_grad", test_custom_overload_grad, devices=devices)
+class TestGradCustoms(unittest.TestCase):
+    pass
 
-    return TestGradCustoms
+
+add_function_test(TestGradCustoms, "test_custom_replay_grad", test_custom_replay_grad, devices=devices)
+add_function_test(TestGradCustoms, "test_custom_overload_grad", test_custom_overload_grad, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2, failfast=False)

warp/tests/test_hash_grid.py

@@ -10,8 +10,7 @@ import unittest
 import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
-
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -73,6 +72,7 @@ def count_neighbors_reference(
 
 
 def test_hashgrid_query(test, device):
+    wp.load_module(device=device)
     rng = np.random.default_rng(123)
 
     grid = wp.HashGrid(dim_x, dim_y, dim_z, device)
@@ -134,19 +134,15 @@ def test_hashgrid_query(test, device):
         test.assertTrue(np.array_equal(counts, counts_ref))
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
 
-    class TestHashGrid(parent):
-        pass
 
-    add_function_test(TestHashGrid, "test_hashgrid_query", test_hashgrid_query, devices=devices)
+class TestHashGrid(unittest.TestCase):
+    pass
 
-    return TestHashGrid
 
+add_function_test(TestHashGrid, "test_hashgrid_query", test_hashgrid_query, devices=devices)
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    wp.force_load()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2, failfast=False)

warp/tests/test_import.py

@@ -5,20 +5,14 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-# include parent path
-import numpy as np
-import math
+import unittest
 
 import warp as wp
-from warp.tests.test_base import *
-
-import unittest
+import warp.tests.test_func as test_func
+from warp.tests.unittest_utils import *
 
 wp.init()
 
-# from test_func import sqr
-import warp.tests.test_func as test_func
-
 
 @wp.kernel
 def test_import_func():
@@ -30,24 +24,16 @@ def test_import_func():
     wp.expect_eq(y, 8.0)
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestImport(parent):
-        pass
+class TestImport(unittest.TestCase):
+    pass
 
-    add_kernel_test(TestImport, kernel=test_import_func, name="test_import_func", dim=1, devices=devices)
 
-    return TestImport
+add_kernel_test(TestImport, kernel=test_import_func, name="test_import_func", dim=1, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    c = register(unittest.TestCase)
-    # unittest.main(verbosity=2)
-
-    wp.force_load()
-
-    loader = unittest.defaultTestLoader
-    testSuite = loader.loadTestsFromTestCase(c)
-    testSuite.debug()
+    unittest.main(verbosity=2)

warp/tests/test_indexedarray.py

@@ -12,7 +12,7 @@ import numpy as np
 
 import warp as wp
 from warp.tests.test_array import FillStruct
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -1106,31 +1106,29 @@ def test_indexedarray_fill_struct(test, device):
     assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
 
-    class TestIndexedArray(parent):
-        pass
 
-    add_function_test(TestIndexedArray, "test_indexedarray_1d", test_indexedarray_1d, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_2d", test_indexedarray_2d, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_3d", test_indexedarray_3d, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_4d", test_indexedarray_4d, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_mixed", test_indexedarray_mixed, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_shape", test_indexedarray_shape, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_getitem", test_indexedarray_getitem, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_slicing", test_indexedarray_slicing, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_generics", test_indexedarray_generics, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_empty", test_indexedarray_empty, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_fill_scalar", test_indexedarray_fill_scalar, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_fill_vector", test_indexedarray_fill_vector, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_fill_matrix", test_indexedarray_fill_matrix, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_fill_struct", test_indexedarray_fill_struct, devices=devices)
+class TestIndexedArray(unittest.TestCase):
+    pass
 
-    return TestIndexedArray
+
+add_function_test(TestIndexedArray, "test_indexedarray_1d", test_indexedarray_1d, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_2d", test_indexedarray_2d, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_3d", test_indexedarray_3d, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_4d", test_indexedarray_4d, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_mixed", test_indexedarray_mixed, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_shape", test_indexedarray_shape, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_getitem", test_indexedarray_getitem, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_slicing", test_indexedarray_slicing, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_generics", test_indexedarray_generics, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_empty", test_indexedarray_empty, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_fill_scalar", test_indexedarray_fill_scalar, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_fill_vector", test_indexedarray_fill_vector, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_fill_matrix", test_indexedarray_fill_matrix, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_fill_struct", test_indexedarray_fill_struct, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_intersect.py

@@ -1,8 +1,16 @@
+# Copyright (c) 2022 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 warp as wp
 import numpy as np
-from warp.tests.test_base import *
+
+import warp as wp
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -44,18 +52,16 @@ def test_intersect_tri(test, device):
     assert_np_equal(result.numpy(), np.array([0]))
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestIntersect(parent):
-        pass
+class TestIntersect(unittest.TestCase):
+    pass
 
-    add_function_test(TestIntersect, "test_intersect_tri", test_intersect_tri, devices=devices)
 
-    return TestIntersect
+add_function_test(TestIntersect, "test_intersect_tri", test_intersect_tri, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2, failfast=False)

warp/tests/test_large.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
+# 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
@@ -9,7 +9,7 @@ import math
 import unittest
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -114,30 +114,28 @@ def test_large_arrays_fast(test, device):
     assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
 
-    class TestLarge(parent):
-        pass
 
-    add_function_test(
-        TestLarge, "test_large_launch_large_kernel", test_large_launch_large_kernel, devices=wp.get_cuda_devices()
-    )
+class TestLarge(unittest.TestCase):
+    pass
 
-    add_function_test(TestLarge, "test_large_launch_max_blocks", test_large_launch_max_blocks, devices=devices)
-    add_function_test(
-        TestLarge,
-        "test_large_launch_very_large_kernel",
-        test_large_launch_very_large_kernel,
-        devices=wp.get_cuda_devices(),
-    )
 
-    add_function_test(TestLarge, "test_large_arrays_fast", test_large_arrays_fast, devices=devices)
+add_function_test(
+    TestLarge, "test_large_launch_large_kernel", test_large_launch_large_kernel, devices=get_unique_cuda_test_devices()
+)
+
+add_function_test(TestLarge, "test_large_launch_max_blocks", test_large_launch_max_blocks, devices=devices)
+add_function_test(
+    TestLarge,
+    "test_large_launch_very_large_kernel",
+    test_large_launch_very_large_kernel,
+    devices=get_unique_cuda_test_devices(),
+)
 
-    return TestLarge
+add_function_test(TestLarge, "test_large_arrays_fast", test_large_arrays_fast, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_launch.py

@@ -7,11 +7,10 @@
 
 import unittest
 
-# include parent path
 import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -306,27 +305,25 @@ def test_launch_tuple_args(test, device):
     assert_np_equal(out.numpy(), np.array((0, 3, 6, 9)))
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
 
-    class TestLaunch(parent):
-        pass
 
-    add_function_test(TestLaunch, "test_launch_1d", test1d, devices=devices)
-    add_function_test(TestLaunch, "test_launch_2d", test2d, devices=devices)
-    add_function_test(TestLaunch, "test_launch_3d", test3d, devices=devices)
-    add_function_test(TestLaunch, "test_launch_4d", test4d, devices=devices)
+class TestLaunch(unittest.TestCase):
+    pass
 
-    add_function_test(TestLaunch, "test_launch_cmd", test_launch_cmd, devices=devices)
-    add_function_test(TestLaunch, "test_launch_cmd_set_param", test_launch_cmd_set_param, devices=devices)
-    add_function_test(TestLaunch, "test_launch_cmd_set_ctype", test_launch_cmd_set_ctype, devices=devices)
-    add_function_test(TestLaunch, "test_launch_cmd_set_dim", test_launch_cmd_set_dim, devices=devices)
-    add_function_test(TestLaunch, "test_launch_cmd_empty", test_launch_cmd_empty, devices=devices)
 
-    return TestLaunch
+add_function_test(TestLaunch, "test_launch_1d", test1d, devices=devices)
+add_function_test(TestLaunch, "test_launch_2d", test2d, devices=devices)
+add_function_test(TestLaunch, "test_launch_3d", test3d, devices=devices)
+add_function_test(TestLaunch, "test_launch_4d", test4d, devices=devices)
+
+add_function_test(TestLaunch, "test_launch_cmd", test_launch_cmd, devices=devices)
+add_function_test(TestLaunch, "test_launch_cmd_set_param", test_launch_cmd_set_param, devices=devices)
+add_function_test(TestLaunch, "test_launch_cmd_set_ctype", test_launch_cmd_set_ctype, devices=devices)
+add_function_test(TestLaunch, "test_launch_cmd_set_dim", test_launch_cmd_set_dim, devices=devices)
+add_function_test(TestLaunch, "test_launch_cmd_empty", test_launch_cmd_empty, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_lerp.py

@@ -5,14 +5,16 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import unittest
 from dataclasses import dataclass
 from typing import Any
-import unittest
 
 import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
+
+wp.init()
 
 
 @dataclass
@@ -162,8 +164,6 @@ TEST_DATA = {
     ),
 }
 
-wp.init()
-
 
 def test_lerp(test, device):
     def make_kernel_fn(data_type):
@@ -184,78 +184,78 @@ def test_lerp(test, device):
             key=f"test_lerp_{data_type.__name__}_kernel",
         )
 
-        for test_data in TEST_DATA[data_type]:
-            a = wp.array(
-                [test_data.a],
-                dtype=data_type,
-                device=device,
-                requires_grad=True,
-            )
-            b = wp.array(
-                [test_data.b],
-                dtype=data_type,
-                device=device,
-                requires_grad=True,
-            )
-            t = wp.array(
-                [test_data.t],
-                dtype=float,
-                device=device,
-                requires_grad=True,
-            )
-            out = wp.array(
-                [0] * wp.types.type_length(data_type),
-                dtype=data_type,
-                device=device,
-                requires_grad=True,
-            )
-
-            tape = wp.Tape()
-            with tape:
-                wp.launch(
-                    kernel,
-                    dim=1,
-                    inputs=[a, b, t, out],
+        with test.subTest(data_type=data_type):
+            for test_data in TEST_DATA[data_type]:
+                a = wp.array(
+                    [test_data.a],
+                    dtype=data_type,
                     device=device,
+                    requires_grad=True,
                 )
-
-            assert_np_equal(
-                out.numpy(),
-                np.array([test_data.expected]),
-                tol=1e-6,
-            )
-
-            if test_data.check_backwards():
-                tape.backward(out)
-
-                assert_np_equal(
-                    tape.gradients[a].numpy(),
-                    np.array([test_data.expected_adj_a]),
-                    tol=1e-6,
+                b = wp.array(
+                    [test_data.b],
+                    dtype=data_type,
+                    device=device,
+                    requires_grad=True,
                 )
-                assert_np_equal(
-                    tape.gradients[b].numpy(),
-                    np.array([test_data.expected_adj_b]),
-                    tol=1e-6,
+                t = wp.array(
+                    [test_data.t],
+                    dtype=float,
+                    device=device,
+                    requires_grad=True,
+                )
+                out = wp.array(
+                    [0] * wp.types.type_length(data_type),
+                    dtype=data_type,
+                    device=device,
+                    requires_grad=True,
                 )
+
+                tape = wp.Tape()
+                with tape:
+                    wp.launch(
+                        kernel,
+                        dim=1,
+                        inputs=[a, b, t, out],
+                        device=device,
+                    )
+
                 assert_np_equal(
-                    tape.gradients[t].numpy(),
-                    np.array([test_data.expected_adj_t]),
+                    out.numpy(),
+                    np.array([test_data.expected]),
                     tol=1e-6,
                 )
 
+                if test_data.check_backwards():
+                    tape.backward(out)
+
+                    assert_np_equal(
+                        tape.gradients[a].numpy(),
+                        np.array([test_data.expected_adj_a]),
+                        tol=1e-6,
+                    )
+                    assert_np_equal(
+                        tape.gradients[b].numpy(),
+                        np.array([test_data.expected_adj_b]),
+                        tol=1e-6,
+                    )
+                    assert_np_equal(
+                        tape.gradients[t].numpy(),
+                        np.array([test_data.expected_adj_t]),
+                        tol=1e-6,
+                    )
+
+
+devices = get_test_devices()
+
 
-def register(parent):
-    devices = get_test_devices()
+class TestLerp(unittest.TestCase):
+    pass
 
-    class TestLerp(parent):
-        pass
 
-    add_function_test(TestLerp, "test_lerp", test_lerp, devices=devices)
-    return TestLerp
+add_function_test(TestLerp, "test_lerp", test_lerp, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)