warp-lang 1.1.0__py3-none-macosx_10_13_universal2.whl → 1.2.0__py3-none-macosx_10_13_universal2.whl

warp/tests/test_import.py

@@ -11,8 +11,6 @@ import warp as wp
 import warp.tests.test_func as test_func
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def test_import_func():

warp/tests/test_indexedarray.py

@@ -14,8 +14,6 @@ import warp as wp
 from warp.tests.test_array import FillStruct
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def kernel_1d(a: wp.indexedarray(dtype=float), expected: wp.array(dtype=float)):

warp/tests/test_intersect.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def intersect_tri(

warp/tests/test_jax.py

@@ -14,8 +14,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 # basic kernel with one input and output
 @wp.kernel

warp/tests/test_large.py

@@ -7,12 +7,11 @@
 
 import math
 import unittest
+from typing import Any
 
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def conditional_sum(result: wp.array(dtype=wp.uint64)):
@@ -86,7 +85,7 @@ def test_large_arrays_slow(test, device):
         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():
+        for wptype in wp.types.scalar_types:
             a1 = wp.zeros(shape_tuple, dtype=wptype, device=device)
             assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
 
@@ -97,22 +96,25 @@ def test_large_arrays_slow(test, device):
             assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
 
 
+@wp.kernel
+def check_array_equal_value(data: wp.array2d(dtype=Any), expect: Any):
+    i, j = wp.tid()
+    wp.expect_eq(data[i, j], expect)
+
+
 def test_large_arrays_fast(test, device):
     # A truncated version of test_large_arrays_slow meant to catch basic errors
 
     # Make is so that a (dim_x, dim_x) array has more than 2**31 elements
     dim_x = math.ceil(math.sqrt(2**31))
 
-    nptype = np.dtype(np.int8)
-    wptype = wp.types.np_dtype_to_warp_type[nptype]
-
-    a1 = wp.zeros((dim_x, dim_x), dtype=wptype, device=device)
+    a1 = wp.zeros((dim_x, dim_x), dtype=wp.int8, device=device)
     a1.fill_(127)
 
-    assert_np_equal(a1.numpy(), 127 * np.ones_like(a1.numpy()))
+    wp.launch(check_array_equal_value, a1.shape, inputs=[a1, wp.int8(127)], device=device)
 
     a1.zero_()
-    assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
+    wp.launch(check_array_equal_value, a1.shape, inputs=[a1, wp.int8(0)], device=device)
 
 
 def test_large_array_excessive_zeros(test, device):

warp/tests/test_launch.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 dim_x = wp.constant(2)
 dim_y = wp.constant(2)
 dim_z = wp.constant(2)

warp/tests/test_lerp.py

@@ -14,8 +14,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @dataclass
 class TestData:
@@ -179,71 +177,29 @@ def test_lerp(test, device):
 
     for data_type in TEST_DATA:
         kernel_fn = make_kernel_fn(data_type)
-        kernel = wp.Kernel(
-            func=kernel_fn,
-            key=f"test_lerp_{data_type.__name__}_kernel",
-        )
+        kernel = wp.Kernel(func=kernel_fn, key=f"test_lerp_{data_type.__name__}_kernel")
 
         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,
-                )
-                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,
-                )
+                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,
+                    [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,
-                    )
+                    wp.launch(kernel, dim=1, inputs=[a, b, t, out], device=device)
 
-                assert_np_equal(
-                    out.numpy(),
-                    np.array([test_data.expected]),
-                    tol=1e-6,
-                )
+                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,
-                    )
-                    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,
-                    )
+                    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()

warp/tests/test_linear_solvers.py

@@ -6,9 +6,7 @@ import warp as wp
 from warp.optim.linear import bicgstab, cg, cr, gmres, preconditioner
 from warp.tests.unittest_utils import *
 
-wp.init()
-
-from warp.context import runtime  # noqa: E402
+wp.init()  # For runtime.core.is_cutlass_enabled()
 
 
 def _check_linear_solve(test, A, b, func, *args, **kwargs):
@@ -172,11 +170,11 @@ class TestLinearSolvers(unittest.TestCase):
 
 devices = get_test_devices()
 
-if not runtime.core.is_cutlass_enabled():
+if not wp.context.runtime.core.is_cutlass_enabled():
     devices = [d for d in devices if not d.is_cuda]
     print("Skipping CUDA linear solver tests because CUTLASS is not supported in this build")
 
-if runtime.core.is_debug_enabled():
+if wp.context.runtime.core.is_debug_enabled():
     # cutlass-based matmul is *very* slow in debug mode -- skip
     devices = [d for d in devices if not d.is_cuda]
     print("Skipping CUDA linear solver tests in debug mode")

warp/tests/test_lvalue.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def rmw_array_kernel(foos: wp.array(dtype=wp.uint32)):

warp/tests/test_marching_cubes.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def make_field(field: wp.array3d(dtype=float), center: wp.vec3, radius: float):

warp/tests/test_mat.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 np_signed_int_types = [
     np.int8,
     np.int16,

warp/tests/test_mat_lite.py

@@ -10,8 +10,6 @@ import unittest
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 mat32d = wp.mat(shape=(3, 2), dtype=wp.float64)
 
 

warp/tests/test_mat_scalar_ops.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 np_signed_int_types = [
     np.int8,
     np.int16,

warp/tests/test_math.py

@@ -13,8 +13,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import add_function_test, assert_np_equal, get_test_devices
 
-wp.init()
-
 
 class ScalarFloatValues(NamedTuple):
     degrees: wp.float32 = None

warp/tests/test_matmul.py

@@ -13,9 +13,7 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
-from warp.context import runtime  # noqa: E402
+wp.init()  # For wp.context.runtime.core.is_cutlass_enabled()
 
 
 class gemm_test_bed_runner:
@@ -81,7 +79,7 @@ class gemm_test_bed_runner:
             tape.backward(grads={D: ones})
 
             D_np = alpha * (A.numpy() @ B.numpy()) + beta * C.numpy()
-            assert np.array_equal(D_np, D.numpy())
+            assert_np_equal(D.numpy(), D_np)
 
             adj_A_np = alpha * np.matmul(ones.numpy(), B.numpy().transpose())
             adj_B_np = alpha * (A.numpy().transpose() @ ones.numpy())
@@ -94,15 +92,15 @@ class gemm_test_bed_runner:
             tape.backward(grads={D: ones})
 
             D_np = alpha * np.matmul(A.numpy(), B.numpy()) + beta * C.numpy()
-            assert np.array_equal(D_np, D.numpy())
+            assert_np_equal(D.numpy(), D_np)
 
             adj_A_np = alpha * np.matmul(ones.numpy(), B.numpy().transpose((0, 2, 1)))
             adj_B_np = alpha * np.matmul(A.numpy().transpose((0, 2, 1)), ones.numpy())
             adj_C_np = beta * ones.numpy()
 
-        assert np.array_equal(adj_A_np, A.grad.numpy())
-        assert np.array_equal(adj_B_np, B.grad.numpy())
-        assert np.array_equal(adj_C_np, C.grad.numpy())
+        assert_np_equal(A.grad.numpy(), adj_A_np)
+        assert_np_equal(B.grad.numpy(), adj_B_np)
+        assert_np_equal(C.grad.numpy(), adj_C_np)
 
     def run(self):
         Ms = [64, 128, 512]
@@ -203,9 +201,9 @@ class gemm_test_bed_runner_transpose:
             tape.backward(grads={D1: ones1, D2: ones2, D3: ones3})
 
             D_np = alpha * (A.numpy() @ B.numpy()) + beta * C1.numpy()
-            assert np.array_equal(D_np, D1.numpy())
-            assert np.array_equal(D_np, D2.numpy())
-            assert np.array_equal(D_np, D3.numpy())
+            assert_np_equal(D1.numpy(), D_np)
+            assert_np_equal(D2.numpy(), D_np)
+            assert_np_equal(D3.numpy(), D_np)
 
             adj_A_np = alpha * (ones1.numpy() @ B.numpy().transpose())
             adj_B_np = alpha * (A.numpy().transpose() @ ones1.numpy())
@@ -224,23 +222,23 @@ class gemm_test_bed_runner_transpose:
             tape.backward(grads={D1: ones1, D2: ones2, D3: ones3})
 
             D_np = alpha * np.matmul(A.numpy(), B.numpy()) + beta * C1.numpy()
-            assert np.array_equal(D_np, D1.numpy())
-            assert np.array_equal(D_np, D2.numpy())
-            assert np.array_equal(D_np, D3.numpy())
+            assert_np_equal(D1.numpy(), D_np)
+            assert_np_equal(D2.numpy(), D_np)
+            assert_np_equal(D3.numpy(), D_np)
 
             adj_A_np = alpha * np.matmul(ones1.numpy(), B.numpy().transpose((0, 2, 1)))
             adj_B_np = alpha * np.matmul(A.numpy().transpose((0, 2, 1)), ones1.numpy())
             adj_C_np = beta * ones1.numpy()
 
-        assert np.array_equal(adj_A_np, A.grad.numpy())
-        assert np.array_equal(adj_A_np, ATT1.grad.numpy())
-        assert np.array_equal(adj_A_np, ATT2.grad.numpy())
-        assert np.array_equal(adj_B_np, B.grad.numpy())
-        assert np.array_equal(adj_B_np, BTT1.grad.numpy())
-        assert np.array_equal(adj_B_np, BTT2.grad.numpy())
-        assert np.array_equal(adj_C_np, C1.grad.numpy())
-        assert np.array_equal(adj_C_np, C2.grad.numpy())
-        assert np.array_equal(adj_C_np, C3.grad.numpy())
+        assert_np_equal(A.grad.numpy(), adj_A_np)
+        assert_np_equal(ATT1.grad.numpy(), adj_A_np)
+        assert_np_equal(ATT2.grad.numpy(), adj_A_np)
+        assert_np_equal(B.grad.numpy(), adj_B_np)
+        assert_np_equal(BTT1.grad.numpy(), adj_B_np)
+        assert_np_equal(BTT2.grad.numpy(), adj_B_np)
+        assert_np_equal(C1.grad.numpy(), adj_C_np)
+        assert_np_equal(C2.grad.numpy(), adj_C_np)
+        assert_np_equal(C3.grad.numpy(), adj_C_np)
 
     def run(self):
         m = 16
@@ -260,13 +258,13 @@ def test_f16(test, device):
     gemm_test_bed_runner_transpose(wp.float16, device).run()
 
 
-@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
+@unittest.skipUnless(wp.context.runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_f32(test, device):
     gemm_test_bed_runner(wp.float32, device).run()
     gemm_test_bed_runner_transpose(wp.float32, device).run()
 
 
-@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
+@unittest.skipUnless(wp.context.runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_f64(test, device):
     gemm_test_bed_runner(wp.float64, device).run()
     gemm_test_bed_runner_transpose(wp.float64, device).run()
@@ -278,7 +276,7 @@ def matrix_sum_kernel(arr: wp.array2d(dtype=float), loss: wp.array(dtype=float))
     wp.atomic_add(loss, 0, arr[i, j])
 
 
-@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
+@unittest.skipUnless(wp.context.runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_tape(test, device):
     rng = np.random.default_rng(42)
     low = -4.5
@@ -318,7 +316,7 @@ def test_tape(test, device):
     assert_array_equal(A.grad, wp.zeros_like(A))
 
 
-@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
+@unittest.skipUnless(wp.context.runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_operator(test, device):
     rng = np.random.default_rng(42)
     low = -4.5
@@ -354,7 +352,7 @@ def test_operator(test, device):
     assert_array_equal(A.grad, wp.zeros_like(A))
 
 
-@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
+@unittest.skipUnless(wp.context.runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_large_batch_count(test, device):
     rng = np.random.default_rng(42)
     low = -4.5
@@ -394,18 +392,18 @@ def test_large_batch_count(test, device):
     tape.backward(grads={D: ones})
 
     D_np = alpha * np.matmul(A.numpy(), B.numpy()) + beta * C.numpy()
-    assert np.array_equal(D_np, D.numpy())
+    assert_np_equal(D.numpy(), D_np)
 
     adj_A_np = alpha * np.matmul(ones.numpy(), B.numpy().transpose((0, 2, 1)))
     adj_B_np = alpha * np.matmul(A.numpy().transpose((0, 2, 1)), ones.numpy())
     adj_C_np = beta * ones.numpy()
 
-    assert np.array_equal(adj_A_np, A.grad.numpy())
-    assert np.array_equal(adj_B_np, B.grad.numpy())
-    assert np.array_equal(adj_C_np, C.grad.numpy())
+    assert_np_equal(A.grad.numpy(), adj_A_np)
+    assert_np_equal(B.grad.numpy(), adj_B_np)
+    assert_np_equal(C.grad.numpy(), adj_C_np)
 
 
-@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
+@unittest.skipUnless(wp.context.runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_adjoint_accumulation(test, device):
     a_np = np.ones(shape=(2, 3))
     b_np = np.ones(shape=(3, 2))
@@ -429,12 +427,12 @@ def test_adjoint_accumulation(test, device):
     grads = {d2_wp: d_grad}
     tape.backward(grads=grads)
 
-    assert np.array_equal(a_wp.grad.numpy(), 4.0 * np.ones(shape=(2, 3)))
-    assert np.array_equal(b_wp.grad.numpy(), 4.0 * np.ones(shape=(3, 2)))
-    assert np.array_equal(c_wp.grad.numpy(), np.ones(shape=(2, 2)))
+    assert_np_equal(a_wp.grad.numpy(), 4.0 * np.ones(shape=(2, 3)))
+    assert_np_equal(b_wp.grad.numpy(), 4.0 * np.ones(shape=(3, 2)))
+    assert_np_equal(c_wp.grad.numpy(), np.ones(shape=(2, 2)))
 
 
-@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
+@unittest.skipUnless(wp.context.runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_cuda_graph_capture(test, device):
     @wp.kernel
     def mat_sum(mat: wp.array2d(dtype=Any), loss: wp.array(dtype=Any)):

warp/tests/test_matmul_lite.py

@@ -12,9 +12,7 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
-from warp.context import runtime  # noqa: E402
+wp.init()  # For wp.context.runtime.core.is_cutlass_enabled()
 
 
 class gemm_test_bed_runner:
@@ -80,7 +78,7 @@ class gemm_test_bed_runner:
             tape.backward(grads={D: ones})
 
             D_np = alpha * (A.numpy() @ B.numpy()) + beta * C.numpy()
-            assert np.array_equal(D_np, D.numpy())
+            assert_np_equal(D.numpy(), D_np)
 
             adj_A_np = alpha * np.matmul(ones.numpy(), B.numpy().transpose())
             adj_B_np = alpha * (A.numpy().transpose() @ ones.numpy())
@@ -93,15 +91,15 @@ class gemm_test_bed_runner:
             tape.backward(grads={D: ones})
 
             D_np = alpha * np.matmul(A.numpy(), B.numpy()) + beta * C.numpy()
-            assert np.array_equal(D_np, D.numpy())
+            assert_np_equal(D.numpy(), D_np)
 
             adj_A_np = alpha * np.matmul(ones.numpy(), B.numpy().transpose((0, 2, 1)))
             adj_B_np = alpha * np.matmul(A.numpy().transpose((0, 2, 1)), ones.numpy())
             adj_C_np = beta * ones.numpy()
 
-        assert np.array_equal(adj_A_np, A.grad.numpy())
-        assert np.array_equal(adj_B_np, B.grad.numpy())
-        assert np.array_equal(adj_C_np, C.grad.numpy())
+        assert_np_equal(A.grad.numpy(), adj_A_np)
+        assert_np_equal(B.grad.numpy(), adj_B_np)
+        assert_np_equal(C.grad.numpy(), adj_C_np)
 
     def run(self):
         Ms = [8]
@@ -202,9 +200,9 @@ class gemm_test_bed_runner_transpose:
             tape.backward(grads={D1: ones1, D2: ones2, D3: ones3})
 
             D_np = alpha * (A.numpy() @ B.numpy()) + beta * C1.numpy()
-            assert np.array_equal(D_np, D1.numpy())
-            assert np.array_equal(D_np, D2.numpy())
-            assert np.array_equal(D_np, D3.numpy())
+            assert_np_equal(D1.numpy(), D_np)
+            assert_np_equal(D2.numpy(), D_np)
+            assert_np_equal(D3.numpy(), D_np)
 
             adj_A_np = alpha * (ones1.numpy() @ B.numpy().transpose())
             adj_B_np = alpha * (A.numpy().transpose() @ ones1.numpy())
@@ -223,23 +221,23 @@ class gemm_test_bed_runner_transpose:
             tape.backward(grads={D1: ones1, D2: ones2, D3: ones3})
 
             D_np = alpha * np.matmul(A.numpy(), B.numpy()) + beta * C1.numpy()
-            assert np.array_equal(D_np, D1.numpy())
-            assert np.array_equal(D_np, D2.numpy())
-            assert np.array_equal(D_np, D3.numpy())
+            assert_np_equal(D1.numpy(), D_np)
+            assert_np_equal(D2.numpy(), D_np)
+            assert_np_equal(D3.numpy(), D_np)
 
             adj_A_np = alpha * np.matmul(ones1.numpy(), B.numpy().transpose((0, 2, 1)))
             adj_B_np = alpha * np.matmul(A.numpy().transpose((0, 2, 1)), ones1.numpy())
             adj_C_np = beta * ones1.numpy()
 
-        assert np.array_equal(adj_A_np, A.grad.numpy())
-        assert np.array_equal(adj_A_np, ATT1.grad.numpy())
-        assert np.array_equal(adj_A_np, ATT2.grad.numpy())
-        assert np.array_equal(adj_B_np, B.grad.numpy())
-        assert np.array_equal(adj_B_np, BTT1.grad.numpy())
-        assert np.array_equal(adj_B_np, BTT2.grad.numpy())
-        assert np.array_equal(adj_C_np, C1.grad.numpy())
-        assert np.array_equal(adj_C_np, C2.grad.numpy())
-        assert np.array_equal(adj_C_np, C3.grad.numpy())
+        assert_np_equal(A.grad.numpy(), adj_A_np)
+        assert_np_equal(ATT1.grad.numpy(), adj_A_np)
+        assert_np_equal(ATT2.grad.numpy(), adj_A_np)
+        assert_np_equal(B.grad.numpy(), adj_B_np)
+        assert_np_equal(BTT1.grad.numpy(), adj_B_np)
+        assert_np_equal(BTT2.grad.numpy(), adj_B_np)
+        assert_np_equal(C1.grad.numpy(), adj_C_np)
+        assert_np_equal(C2.grad.numpy(), adj_C_np)
+        assert_np_equal(C3.grad.numpy(), adj_C_np)
 
     def run(self):
         m = 8
@@ -253,7 +251,7 @@ class gemm_test_bed_runner_transpose:
             self.run_and_verify(m, n, k, batch_count, alpha, beta)
 
 
-@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
+@unittest.skipUnless(wp.context.runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_f32(test, device):
     gemm_test_bed_runner(wp.float32, device).run()
     gemm_test_bed_runner_transpose(wp.float32, device).run()
@@ -265,7 +263,7 @@ def matrix_sum_kernel(arr: wp.array2d(dtype=float), loss: wp.array(dtype=float))
     wp.atomic_add(loss, 0, arr[i, j])
 
 
-@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
+@unittest.skipUnless(wp.context.runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_tape(test, device):
     rng = np.random.default_rng(42)
     low = -4.5
@@ -305,7 +303,7 @@ def test_tape(test, device):
     assert_array_equal(A.grad, wp.zeros_like(A))
 
 
-@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
+@unittest.skipUnless(wp.context.runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_operator(test, device):
     rng = np.random.default_rng(42)
     low = -4.5
@@ -341,7 +339,7 @@ def test_operator(test, device):
     assert_array_equal(A.grad, wp.zeros_like(A))
 
 
-@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
+@unittest.skipUnless(wp.context.runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_large_batch_count(test, device):
     rng = np.random.default_rng(42)
     low = -4.5
@@ -381,15 +379,15 @@ def test_large_batch_count(test, device):
     tape.backward(grads={D: ones})
 
     D_np = alpha * np.matmul(A.numpy(), B.numpy()) + beta * C.numpy()
-    assert np.array_equal(D_np, D.numpy())
+    assert_np_equal(D.numpy(), D_np)
 
     adj_A_np = alpha * np.matmul(ones.numpy(), B.numpy().transpose((0, 2, 1)))
     adj_B_np = alpha * np.matmul(A.numpy().transpose((0, 2, 1)), ones.numpy())
     adj_C_np = beta * ones.numpy()
 
-    assert np.array_equal(adj_A_np, A.grad.numpy())
-    assert np.array_equal(adj_B_np, B.grad.numpy())
-    assert np.array_equal(adj_C_np, C.grad.numpy())
+    assert_np_equal(A.grad.numpy(), adj_A_np)
+    assert_np_equal(B.grad.numpy(), adj_B_np)
+    assert_np_equal(C.grad.numpy(), adj_C_np)
 
 
 devices = get_test_devices()

warp/tests/test_mempool.py

@@ -10,8 +10,6 @@ import unittest
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 def get_device_pair_with_mempool_access_support():
     devices = wp.get_cuda_devices()

warp/tests/test_mesh.py

@@ -65,9 +65,6 @@ VERTEX_COUNT = 36
 FACE_COUNT = 12
 
 
-wp.init()
-
-
 @wp.kernel(enable_backward=False)
 def read_points_kernel(
     mesh_id: wp.uint64,

warp/tests/test_mesh_query_aabb.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.func
 def min_vec3(a: wp.vec3, b: wp.vec3):

warp/tests/test_mesh_query_point.py

@@ -13,8 +13,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def sample_mesh_query(

warp/tests/test_mesh_query_ray.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 # triangulate a list of polygon face indices
 def triangulate(face_counts, face_indices):

warp/tests/test_mlp.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.func
 def mlp_activation(z: float):

warp/tests/test_model.py

@@ -13,8 +13,6 @@ import warp as wp
 from warp.sim import ModelBuilder
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 class TestModel(unittest.TestCase):
     def test_add_triangles(self):