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

warp/tests/test_rand.py

@@ -5,12 +5,15 @@
 # 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.tests.unittest_utils import *
+
 # import matplotlib.pyplot as plt
 
-import warp as wp
-from warp.tests.test_base import *
 
 wp.init()
 
@@ -60,7 +63,7 @@ def test_rand(test, device):
     wp.copy(int_ab_host, int_ab_device)
     wp.copy(float_01_host, float_01_device)
     wp.copy(float_ab_host, float_ab_device)
-    wp.synchronize()
+    wp.synchronize_device(device)
 
     int_a = int_a_host.numpy()
     int_ab = int_ab_host.numpy()
@@ -243,10 +246,10 @@ def test_poisson(test, device):
     # _ = plt.hist(poisson_high.numpy(), bins)
     # plt.show()
 
-    np.random.default_rng(seed)
+    rng = np.random.default_rng(seed)
 
-    np_poisson_low = np.random.poisson(3.0, N)
-    np_poisson_high = np.random.poisson(42.0, N)
+    np_poisson_low = rng.poisson(lam=3.0, size=N)
+    np_poisson_high = rng.poisson(lam=42.0, size=N)
 
     poisson_low_mean = np.mean(poisson_low.numpy())
     np_poisson_low_mean = np.mean(np_poisson_low)
@@ -267,20 +270,19 @@ def test_poisson(test, device):
     test.assertTrue(np.abs(poisson_high_std - np_poisson_high_std) <= 2e-1)
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestNoise(parent):
-        pass
+class TestRand(unittest.TestCase):
+    pass
 
-    add_function_test(TestNoise, "test_rand", test_rand, devices=devices)
-    add_function_test(TestNoise, "test_sample_cdf", test_sample_cdf, devices=devices)
-    add_function_test(TestNoise, "test_sampling_methods", test_sampling_methods, devices=devices)
-    add_function_test(TestNoise, "test_poisson", test_poisson, devices=devices)
 
-    return TestNoise
+add_function_test(TestRand, "test_rand", test_rand, devices=devices)
+add_function_test(TestRand, "test_sample_cdf", test_sample_cdf, devices=devices)
+add_function_test(TestRand, "test_sampling_methods", test_sampling_methods, devices=devices)
+add_function_test(TestRand, "test_poisson", test_poisson, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_reload.py

@@ -5,29 +5,34 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import numpy as np
-import warp as wp
+import importlib
+import os
+import unittest
 
-import math
+import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
 
-import unittest
-import importlib
-import os
+# dummy modules used for testing reload with dependencies
+import warp.tests.aux_test_dependent as test_dependent
+import warp.tests.aux_test_reference as test_reference
+import warp.tests.aux_test_reference_reference as test_reference_reference
 
 # dummy module used for testing reload
-import warp.tests.test_square as test_square
-
-# dummy modules used for testing reload with dependencies
-import warp.tests.test_dependent as test_dependent
-import warp.tests.test_reference as test_reference
-import warp.tests.test_reference_reference as test_reference_reference
+import warp.tests.aux_test_square as test_square
+from warp.tests.unittest_utils import *
 
 wp.init()
 
 
+def reload_module(module):
+    # Clearing the .pyc file associated with a module is a necessary workaround
+    # for `importlib.reload` to work as expected when run from within Kit.
+    cache_file = importlib.util.cache_from_source(module.__file__)
+    os.remove(cache_file)
+    importlib.reload(module)
+
+
 def test_redefine(test, device):
     # --------------------------------------------
     # first pass
@@ -102,32 +107,32 @@ def run(expect, device):
 
 def test_reload(test, device):
     # write out the module python and import it
-    f = open(os.path.abspath(os.path.join(os.path.dirname(__file__), "test_square.py")), "w")
+    f = open(os.path.abspath(os.path.join(os.path.dirname(__file__), "aux_test_square.py")), "w")
     f.writelines(square_two)
     f.flush()
     f.close()
 
-    importlib.reload(test_square)
+    reload_module(test_square)
     test_square.run(expect=4.0, device=device)  # 2*2=4
 
-    f = open(os.path.abspath(os.path.join(os.path.dirname(__file__), "test_square.py")), "w")
+    f = open(os.path.abspath(os.path.join(os.path.dirname(__file__), "aux_test_square.py")), "w")
     f.writelines(square_four)
     f.flush()
     f.close()
 
     # reload module, this should trigger all of the funcs / kernels to be updated
-    importlib.reload(test_square)
+    reload_module(test_square)
     test_square.run(expect=16.0, device=device)  # 4*4 = 16
 
 
 def test_reload_class(test, device):
     def test_func():
-        import warp.tests.test_class_kernel
-        from warp.tests.test_class_kernel import ClassKernelTest
-
         import importlib as imp
 
-        imp.reload(warp.tests.test_class_kernel)
+        import warp.tests.aux_test_class_kernel
+        from warp.tests.aux_test_class_kernel import ClassKernelTest
+
+        imp.reload(warp.tests.aux_test_class_kernel)
 
         ctest = ClassKernelTest(device)
         expected = np.zeros((10, 3, 3), dtype=np.float32)
@@ -141,7 +146,7 @@ def test_reload_class(test, device):
 template_ref = """# This file is used to test reloading module references.
 
 import warp as wp
-import warp.tests.test_reference_reference as refref
+import warp.tests.aux_test_reference_reference as refref
 
 wp.init()
 
@@ -165,8 +170,8 @@ def more_magic():
 
 
 def test_reload_references(test, device):
-    path_ref = os.path.abspath(os.path.join(os.path.dirname(__file__), "test_reference.py"))
-    path_refref = os.path.abspath(os.path.join(os.path.dirname(__file__), "test_reference_reference.py"))
+    path_ref = os.path.abspath(os.path.join(os.path.dirname(__file__), "aux_test_reference.py"))
+    path_refref = os.path.abspath(os.path.join(os.path.dirname(__file__), "aux_test_reference_reference.py"))
 
     # rewrite both dependency modules and reload them
     with open(path_ref, "w") as f:
@@ -194,20 +199,19 @@ def test_reload_references(test, device):
     test_dependent.run(expect=4.0, device=device)  # 2 * 2 = 4
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestReload(parent):
-        pass
+class TestReload(unittest.TestCase):
+    pass
 
-    add_function_test(TestReload, "test_redefine", test_redefine, devices=devices)
-    add_function_test(TestReload, "test_reload", test_reload, devices=devices)
-    add_function_test(TestReload, "test_reload_class", test_reload_class, devices=devices)
-    add_function_test(TestReload, "test_reload_references", test_reload_references, devices=devices)
 
-    return TestReload
+add_function_test(TestReload, "test_redefine", test_redefine, devices=devices)
+add_function_test(TestReload, "test_reload", test_reload, devices=devices)
+add_function_test(TestReload, "test_reload_class", test_reload_class, devices=devices)
+add_function_test(TestReload, "test_reload_references", test_reload_references, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2, failfast=False)

warp/tests/test_rounding.py

@@ -5,11 +5,13 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import warp as wp
-from warp.tests.test_base import *
+import unittest
 
 import numpy as np
 
+import warp as wp
+from warp.tests.unittest_utils import *
+
 compare_to_numpy = False
 print_results = False
 
@@ -25,6 +27,7 @@ def test_kernel(
     x_cast: wp.array(dtype=float),
     x_floor: wp.array(dtype=float),
     x_ceil: wp.array(dtype=float),
+    x_frac: wp.array(dtype=float),
 ):
     tid = wp.tid()
 
@@ -34,6 +37,7 @@ def test_kernel(
     x_cast[tid] = float(int(x[tid]))
     x_floor[tid] = wp.floor(x[tid])
     x_ceil[tid] = wp.ceil(x[tid])
+    x_frac[tid] = wp.frac(x[tid])
 
 
 def test_rounding(test, device):
@@ -82,8 +86,11 @@ def test_rounding(test, device):
     x_cast = wp.empty(N, dtype=float, device=device)
     x_floor = wp.empty(N, dtype=float, device=device)
     x_ceil = wp.empty(N, dtype=float, device=device)
+    x_frac = wp.empty(N, dtype=float, device=device)
 
-    wp.launch(kernel=test_kernel, dim=N, inputs=[x, x_round, x_rint, x_trunc, x_cast, x_floor, x_ceil], device=device)
+    wp.launch(
+        kernel=test_kernel, dim=N, inputs=[x, x_round, x_rint, x_trunc, x_cast, x_floor, x_ceil, x_frac], device=device
+    )
 
     wp.synchronize()
 
@@ -93,46 +100,47 @@ def test_rounding(test, device):
     nx_cast = x_cast.numpy().reshape(N)
     nx_floor = x_floor.numpy().reshape(N)
     nx_ceil = x_ceil.numpy().reshape(N)
+    nx_frac = x_frac.numpy().reshape(N)
 
-    tab = np.stack([nx, nx_round, nx_rint, nx_trunc, nx_cast, nx_floor, nx_ceil], axis=1)
+    tab = np.stack([nx, nx_round, nx_rint, nx_trunc, nx_cast, nx_floor, nx_ceil, nx_frac], axis=1)
 
     golden = np.array(
         [
-            [4.9, 5.0, 5.0, 4.0, 4.0, 4.0, 5.0],
-            [4.5, 5.0, 4.0, 4.0, 4.0, 4.0, 5.0],
-            [4.1, 4.0, 4.0, 4.0, 4.0, 4.0, 5.0],
-            [3.9, 4.0, 4.0, 3.0, 3.0, 3.0, 4.0],
-            [3.5, 4.0, 4.0, 3.0, 3.0, 3.0, 4.0],
-            [3.1, 3.0, 3.0, 3.0, 3.0, 3.0, 4.0],
-            [2.9, 3.0, 3.0, 2.0, 2.0, 2.0, 3.0],
-            [2.5, 3.0, 2.0, 2.0, 2.0, 2.0, 3.0],
-            [2.1, 2.0, 2.0, 2.0, 2.0, 2.0, 3.0],
-            [1.9, 2.0, 2.0, 1.0, 1.0, 1.0, 2.0],
-            [1.5, 2.0, 2.0, 1.0, 1.0, 1.0, 2.0],
-            [1.1, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0],
-            [0.9, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0],
-            [0.5, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0],
-            [0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
-            [-0.1, -0.0, -0.0, -0.0, 0.0, -1.0, -0.0],
-            [-0.5, -1.0, -0.0, -0.0, 0.0, -1.0, -0.0],
-            [-0.9, -1.0, -1.0, -0.0, 0.0, -1.0, -0.0],
-            [-1.1, -1.0, -1.0, -1.0, -1.0, -2.0, -1.0],
-            [-1.5, -2.0, -2.0, -1.0, -1.0, -2.0, -1.0],
-            [-1.9, -2.0, -2.0, -1.0, -1.0, -2.0, -1.0],
-            [-2.1, -2.0, -2.0, -2.0, -2.0, -3.0, -2.0],
-            [-2.5, -3.0, -2.0, -2.0, -2.0, -3.0, -2.0],
-            [-2.9, -3.0, -3.0, -2.0, -2.0, -3.0, -2.0],
-            [-3.1, -3.0, -3.0, -3.0, -3.0, -4.0, -3.0],
-            [-3.5, -4.0, -4.0, -3.0, -3.0, -4.0, -3.0],
-            [-3.9, -4.0, -4.0, -3.0, -3.0, -4.0, -3.0],
-            [-4.1, -4.0, -4.0, -4.0, -4.0, -5.0, -4.0],
-            [-4.5, -5.0, -4.0, -4.0, -4.0, -5.0, -4.0],
-            [-4.9, -5.0, -5.0, -4.0, -4.0, -5.0, -4.0],
+            [4.9, 5.0, 5.0, 4.0, 4.0, 4.0, 5.0, 0.9],
+            [4.5, 5.0, 4.0, 4.0, 4.0, 4.0, 5.0, 0.5],
+            [4.1, 4.0, 4.0, 4.0, 4.0, 4.0, 5.0, 0.1],
+            [3.9, 4.0, 4.0, 3.0, 3.0, 3.0, 4.0, 0.9],
+            [3.5, 4.0, 4.0, 3.0, 3.0, 3.0, 4.0, 0.5],
+            [3.1, 3.0, 3.0, 3.0, 3.0, 3.0, 4.0, 0.1],
+            [2.9, 3.0, 3.0, 2.0, 2.0, 2.0, 3.0, 0.9],
+            [2.5, 3.0, 2.0, 2.0, 2.0, 2.0, 3.0, 0.5],
+            [2.1, 2.0, 2.0, 2.0, 2.0, 2.0, 3.0, 0.1],
+            [1.9, 2.0, 2.0, 1.0, 1.0, 1.0, 2.0, 0.9],
+            [1.5, 2.0, 2.0, 1.0, 1.0, 1.0, 2.0, 0.5],
+            [1.1, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 0.1],
+            [0.9, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.9],
+            [0.5, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.5],
+            [0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.1],
+            [-0.1, -0.0, -0.0, -0.0, 0.0, -1.0, -0.0, -0.1],
+            [-0.5, -1.0, -0.0, -0.0, 0.0, -1.0, -0.0, -0.5],
+            [-0.9, -1.0, -1.0, -0.0, 0.0, -1.0, -0.0, -0.9],
+            [-1.1, -1.0, -1.0, -1.0, -1.0, -2.0, -1.0, -0.1],
+            [-1.5, -2.0, -2.0, -1.0, -1.0, -2.0, -1.0, -0.5],
+            [-1.9, -2.0, -2.0, -1.0, -1.0, -2.0, -1.0, -0.9],
+            [-2.1, -2.0, -2.0, -2.0, -2.0, -3.0, -2.0, -0.1],
+            [-2.5, -3.0, -2.0, -2.0, -2.0, -3.0, -2.0, -0.5],
+            [-2.9, -3.0, -3.0, -2.0, -2.0, -3.0, -2.0, -0.9],
+            [-3.1, -3.0, -3.0, -3.0, -3.0, -4.0, -3.0, -0.1],
+            [-3.5, -4.0, -4.0, -3.0, -3.0, -4.0, -3.0, -0.5],
+            [-3.9, -4.0, -4.0, -3.0, -3.0, -4.0, -3.0, -0.9],
+            [-4.1, -4.0, -4.0, -4.0, -4.0, -5.0, -4.0, -0.1],
+            [-4.5, -5.0, -4.0, -4.0, -4.0, -5.0, -4.0, -0.5],
+            [-4.9, -5.0, -5.0, -4.0, -4.0, -5.0, -4.0, -0.9],
         ],
         dtype=np.float32,
     )
 
-    assert_np_equal(tab, golden)
+    assert_np_equal(tab, golden, tol=1e-6)
 
     if print_results:
         np.set_printoptions(formatter={"float": lambda x: "{:6.1f}".format(x).replace(".0", ".")})
@@ -149,24 +157,23 @@ def test_rounding(test, device):
         nx_fix = np.fix(nx)
         nx_floor = np.floor(nx)
         nx_ceil = np.ceil(nx)
+        nx_frac = np.modf(nx)[0]
 
-        tab = np.stack([nx, nx_round, nx_rint, nx_trunc, nx_fix, nx_floor, nx_ceil], axis=1)
+        tab = np.stack([nx, nx_round, nx_rint, nx_trunc, nx_fix, nx_floor, nx_ceil, nx_frac], axis=1)
         print("   %5s %5s %5s %5s %5s %5s %5s" % ("x ", "round", "rint", "trunc", "fix", "floor", "ceil"))
         print(tab)
         print("----------------------------------------------")
 
 
-def register(parent):
-    class TestRounding(parent):
-        pass
+class TestRounding(unittest.TestCase):
+    pass
 
-    devices = get_test_devices()
 
-    add_function_test(TestRounding, "test_rounding", test_rounding, devices=devices)
+devices = get_test_devices()
 
-    return TestRounding
+add_function_test(TestRounding, "test_rounding", test_rounding, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_runlength_encode.py

@@ -0,0 +1,190 @@
+# 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
+from functools import partial
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+from warp.utils import runlength_encode
+
+wp.init()
+
+
+def test_runlength_encode_int(test, device, n):
+    rng = np.random.default_rng(123)
+
+    values_np = np.sort(rng.integers(-10, high=10, size=n, dtype=int))
+
+    unique_values_np, unique_counts_np = np.unique(values_np, return_counts=True)
+
+    values = wp.array(values_np, device=device, dtype=int)
+
+    unique_values = wp.empty_like(values)
+    unique_counts = wp.empty_like(values)
+
+    run_count = runlength_encode(values, unique_values, unique_counts)
+
+    test.assertEqual(run_count, len(unique_values_np))
+    assert_np_equal(unique_values.numpy()[:run_count], unique_values_np[:run_count])
+    assert_np_equal(unique_counts.numpy()[:run_count], unique_counts_np[:run_count])
+
+
+def test_runlength_encode_error_insufficient_storage(test, device):
+    values = wp.zeros(123, dtype=int, device=device)
+    run_values = wp.empty(1, dtype=int, device=device)
+    run_lengths = wp.empty(123, dtype=int, device=device)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"Output array storage sizes must be at least equal to value_count$",
+    ):
+        runlength_encode(values, run_values, run_lengths)
+
+    values = wp.zeros(123, dtype=int, device="cpu")
+    run_values = wp.empty(123, dtype=int, device="cpu")
+    run_lengths = wp.empty(1, dtype=int, device="cpu")
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"Output array storage sizes must be at least equal to value_count$",
+    ):
+        runlength_encode(values, run_values, run_lengths)
+
+
+def test_runlength_encode_error_dtypes_mismatch(test, device):
+    values = wp.zeros(123, dtype=int, device=device)
+    run_values = wp.empty(123, dtype=float, device=device)
+    run_lengths = wp.empty_like(values, device=device)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"values and run_values data types do not match$",
+    ):
+        runlength_encode(values, run_values, run_lengths)
+
+
+def test_runlength_encode_error_run_length_unsupported_dtype(test, device):
+    values = wp.zeros(123, dtype=int, device=device)
+    run_values = wp.empty(123, dtype=int, device=device)
+    run_lengths = wp.empty(123, dtype=float, device=device)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"run_lengths array must be of type int32$",
+    ):
+        runlength_encode(values, run_values, run_lengths)
+
+
+def test_runlength_encode_error_run_count_unsupported_dtype(test, device):
+    values = wp.zeros(123, dtype=int, device=device)
+    run_values = wp.empty_like(values, device=device)
+    run_lengths = wp.empty_like(values, device=device)
+    run_count = wp.empty(shape=(1,), dtype=float, device=device)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"run_count array must be of type int32$",
+    ):
+        runlength_encode(values, run_values, run_lengths, run_count=run_count)
+
+
+def test_runlength_encode_error_unsupported_dtype(test, device):
+    values = wp.zeros(123, dtype=float, device=device)
+    run_values = wp.empty(123, dtype=float, device=device)
+    run_lengths = wp.empty(123, dtype=int, device=device)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"Unsupported data type$",
+    ):
+        runlength_encode(values, run_values, run_lengths)
+
+
+devices = get_test_devices()
+
+
+class TestRunlengthEncode(unittest.TestCase):
+    @unittest.skipUnless(wp.is_cuda_available(), "Requires CUDA")
+    def test_runlength_encode_error_devices_mismatch(self):
+        values = wp.zeros(123, dtype=int, device="cpu")
+        run_values = wp.empty_like(values, device="cuda:0")
+        run_lengths = wp.empty_like(values, device="cuda:0")
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"Array storage devices do not match$",
+        ):
+            runlength_encode(values, run_values, run_lengths)
+
+        values = wp.zeros(123, dtype=int, device="cpu")
+        run_values = wp.empty_like(values, device="cpu")
+        run_lengths = wp.empty_like(values, device="cuda:0")
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"Array storage devices do not match$",
+        ):
+            runlength_encode(values, run_values, run_lengths)
+
+        values = wp.zeros(123, dtype=int, device="cpu")
+        run_values = wp.empty_like(values, device="cuda:0")
+        run_lengths = wp.empty_like(values, device="cpu")
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"Array storage devices do not match$",
+        ):
+            runlength_encode(values, run_values, run_lengths)
+
+    @unittest.skipUnless(wp.is_cuda_available(), "Requires CUDA")
+    def test_runlength_encode_error_run_count_device_mismatch(self):
+        values = wp.zeros(123, dtype=int, device="cpu")
+        run_values = wp.empty_like(values, device="cpu")
+        run_lengths = wp.empty_like(values, device="cpu")
+        run_count = wp.empty(shape=(1,), dtype=int, device="cuda:0")
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"run_count storage device does not match other arrays$",
+        ):
+            runlength_encode(values, run_values, run_lengths, run_count=run_count)
+
+
+add_function_test(
+    TestRunlengthEncode, "test_runlength_encode_int", partial(test_runlength_encode_int, n=100), devices=devices
+)
+add_function_test(
+    TestRunlengthEncode, "test_runlength_encode_empty", partial(test_runlength_encode_int, n=0), devices=devices
+)
+add_function_test(
+    TestRunlengthEncode,
+    "test_runlength_encode_error_insufficient_storage",
+    test_runlength_encode_error_insufficient_storage,
+    devices=devices,
+)
+add_function_test(
+    TestRunlengthEncode,
+    "test_runlength_encode_error_dtypes_mismatch",
+    test_runlength_encode_error_dtypes_mismatch,
+    devices=devices,
+)
+add_function_test(
+    TestRunlengthEncode,
+    "test_runlength_encode_error_run_length_unsupported_dtype",
+    test_runlength_encode_error_run_length_unsupported_dtype,
+    devices=devices,
+)
+add_function_test(
+    TestRunlengthEncode,
+    "test_runlength_encode_error_run_count_unsupported_dtype",
+    test_runlength_encode_error_run_count_unsupported_dtype,
+    devices=devices,
+)
+add_function_test(
+    TestRunlengthEncode,
+    "test_runlength_encode_error_unsupported_dtype",
+    test_runlength_encode_error_unsupported_dtype,
+    devices=devices,
+)
+
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_smoothstep.py

@@ -5,14 +5,14 @@
 # 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 *
 
 
 @dataclass
@@ -87,11 +87,9 @@ def test_smoothstep(test, device):
 
     for data_type in TEST_DATA:
         kernel_fn = make_kernel_fn(data_type)
-        module = wp.get_module(kernel_fn.__module__)
         kernel = wp.Kernel(
             func=kernel_fn,
             key=f"test_smoothstep{data_type.__name__}_kernel",
-            module=module,
         )
 
         for test_data in TEST_DATA[data_type]:
@@ -155,16 +153,16 @@ def test_smoothstep(test, device):
                 )
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
+
+class TestSmoothstep(unittest.TestCase):
+    pass
 
-    class TestSmoothstep(parent):
-        pass
 
-    add_function_test(TestSmoothstep, "test_smoothstep", test_smoothstep, devices=devices)
-    return TestSmoothstep
+add_function_test(TestSmoothstep, "test_smoothstep", test_smoothstep, devices=devices)
 
 
 if __name__ == "__main__":
-    _ = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)