warp-lang 1.9.0__py3-none-win_amd64.whl → 1.10.0rc2__py3-none-win_amd64.whl

warp/tests/__main__.py

@@ -1,4 +1,4 @@
-from warp.thirdparty.unittest_parallel import main
+from warp._src.thirdparty.unittest_parallel import main
 
 if __name__ == "__main__":
     main()

warp/tests/cuda/test_async.py

@@ -18,8 +18,8 @@ import unittest
 import numpy as np
 
 import warp as wp
+from warp._src.utils import check_p2p
 from warp.tests.unittest_utils import *
-from warp.utils import check_p2p
 
 
 class Capturable:
@@ -426,13 +426,13 @@ def copy_template(test, src_ctor, dst_ctor, src_device, dst_device, n, params: C
         if expected_error_type is not None:
             # disable error output from Warp if we expect an exception
             try:
-                saved_error_output_enabled = wp.context.runtime.core.wp_is_error_output_enabled()
-                wp.context.runtime.core.wp_set_error_output_enabled(False)
+                saved_error_output_enabled = wp._src.context.runtime.core.wp_is_error_output_enabled()
+                wp._src.context.runtime.core.wp_set_error_output_enabled(False)
                 with test.assertRaisesRegex(expected_error_type, expected_error_regex):
                     with Capturable(use_graph=params.use_graph, stream=stream):
                         wp.copy(dst, src, stream=stream_arg)
             finally:
-                wp.context.runtime.core.wp_set_error_output_enabled(saved_error_output_enabled)
+                wp._src.context.runtime.core.wp_set_error_output_enabled(saved_error_output_enabled)
                 wp.synchronize()
 
                 # print(f"SUCCESSFUL ERROR PREDICTION: {expected_error_regex}")

warp/tests/cuda/test_conditional_captures.py

@@ -20,7 +20,7 @@ import unittest
 import numpy as np
 
 import warp as wp
-from warp.context import assert_conditional_graph_support
+from warp._src.context import assert_conditional_graph_support
 from warp.tests.unittest_utils import *
 
 

warp/tests/cuda/test_multigpu.py

@@ -18,8 +18,8 @@ import unittest
 import numpy as np
 
 import warp as wp
+from warp._src.utils import check_p2p
 from warp.tests.unittest_utils import *
-from warp.utils import check_p2p
 
 
 @wp.kernel

warp/tests/cuda/test_streams.py

@@ -18,8 +18,8 @@ import unittest
 import numpy as np
 
 import warp as wp
+from warp._src.utils import check_p2p
 from warp.tests.unittest_utils import *
-from warp.utils import check_p2p
 
 
 @wp.kernel
@@ -365,6 +365,41 @@ def test_event_elapsed_time_graph(test, device):
     test.assertGreater(elapsed, 0)
 
 
+def test_event_external(test, device):
+    with wp.ScopedDevice(device):
+        # event used to synchronize two graphs (external event)
+        event = wp.Event()
+
+        n = 1_000_000
+        a = wp.zeros(n, dtype=float)
+        b = wp.zeros(n, dtype=float)
+        c = wp.zeros(n, dtype=float)
+
+        with wp.ScopedCapture() as capture1:
+            wp.launch(inc, dim=n, inputs=[a])
+            wp.launch(inc, dim=n, inputs=[b])
+            # record in first graph
+            wp.record_event(event, external=True)
+
+        with wp.ScopedCapture() as capture2:
+            # wait in second graph
+            wp.wait_event(event, external=True)
+            wp.launch(sum, dim=n, inputs=[a, b, c])
+
+        stream1 = wp.Stream()
+        stream2 = wp.Stream()
+        num_iters = 10
+
+        for _ in range(num_iters):
+            # Launch graphs on different streams, but they should be
+            # synchronized using the external event.
+            wp.capture_launch(capture1.graph, stream=stream1)
+            wp.capture_launch(capture2.graph, stream=stream2)
+
+        expected = np.full(n, 2 * num_iters, dtype=np.float32)
+        assert_np_equal(c.numpy(), expected)
+
+
 def test_stream_priority_basics(test, device):
     standard_stream = wp.Stream(device)
     test.assertEqual(standard_stream.priority, 0, "Default priority of streams must be 0.")
@@ -470,6 +505,25 @@ def test_stream_event_is_complete(test, device):
                 test.assertEqual(a.numpy()[0], (iter + 1) * threads)
 
 
+def test_graph_destroy_during_capture(test, device):
+    with wp.ScopedDevice(device):
+        n = 10
+        a = wp.zeros(n, dtype=float)
+
+        with wp.ScopedCapture() as capture1:
+            wp.launch(inc, dim=n, inputs=[a])
+
+        wp.capture_launch(capture1.graph)
+
+        with wp.ScopedCapture() as capture2:
+            del capture1  # <--- should be deferred
+            wp.launch(inc, dim=n, inputs=[a])
+
+        wp.capture_launch(capture2.graph)
+
+        assert_np_equal(a.numpy(), np.full(n, 2, dtype=np.float32))
+
+
 devices = get_selected_cuda_test_devices()
 
 
@@ -628,6 +682,9 @@ add_function_test(TestStreams, "test_stream_event_is_complete", test_stream_even
 add_function_test(TestStreams, "test_event_synchronize", test_event_synchronize, devices=devices)
 add_function_test(TestStreams, "test_event_elapsed_time", test_event_elapsed_time, devices=devices)
 add_function_test(TestStreams, "test_event_elapsed_time_graph", test_event_elapsed_time_graph, devices=devices)
+add_function_test(TestStreams, "test_event_external", test_event_external, devices=devices)
+
+add_function_test(TestStreams, "test_graph_destroy_during_capture", test_graph_destroy_during_capture, devices=devices)
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/geometry/test_bvh.py

@@ -75,7 +75,7 @@ def intersect_ray_aabb(start, rcp_dir, lower, upper):
         return 0
 
 
-def test_bvh(test, type, device):
+def test_bvh(test, type, device, leaf_size):
     rng = np.random.default_rng(123)
 
     num_bounds = 100
@@ -85,7 +85,7 @@ def test_bvh(test, type, device):
     device_lowers = wp.array(lowers, dtype=wp.vec3, device=device)
     device_uppers = wp.array(uppers, dtype=wp.vec3, device=device)
 
-    bvh = wp.Bvh(device_lowers, device_uppers)
+    bvh = wp.Bvh(device_lowers, device_uppers, leaf_size=leaf_size)
 
     bounds_intersected = wp.zeros(shape=(num_bounds), dtype=int, device=device)
 
@@ -95,18 +95,16 @@ def test_bvh(test, type, device):
     query_start = wp.vec3(0.0, 0.0, 0.0)
     query_dir = wp.normalize(wp.vec3(1.0, 1.0, 1.0))
 
-    for test_case in range(2):
+    for test_case in range(3):
         if type == "AABB":
             wp.launch(
-                kernel=bvh_query_aabb,
+                bvh_query_aabb,
                 dim=1,
                 inputs=[bvh.id, query_lower, query_upper, bounds_intersected],
                 device=device,
             )
         else:
-            wp.launch(
-                kernel=bvh_query_ray, dim=1, inputs=[bvh.id, query_start, query_dir, bounds_intersected], device=device
-            )
+            wp.launch(bvh_query_ray, dim=1, inputs=[bvh.id, query_start, query_dir, bounds_intersected], device=device)
 
         device_intersected = bounds_intersected.numpy()
 
@@ -120,48 +118,178 @@ def test_bvh(test, type, device):
 
             test.assertEqual(host_intersected, device_intersected[i])
 
-        if test_case == 0:
+        if test_case == 0 or test_case == 1:
             lowers = rng.random(size=(num_bounds, 3)) * 5.0
             uppers = lowers + rng.random(size=(num_bounds, 3)) * 5.0
             wp.copy(device_lowers, wp.array(lowers, dtype=wp.vec3, device=device))
             wp.copy(device_uppers, wp.array(uppers, dtype=wp.vec3, device=device))
-            bvh.refit()
             bounds_intersected.zero_()
 
+            if test_case == 0:
+                bvh.refit()
+            else:
+                bvh.rebuild()
+
 
 def test_bvh_query_aabb(test, device):
-    test_bvh(test, "AABB", device)
+    for leaf_size in [1, 2, 4]:
+        test_bvh(test, "AABB", device, leaf_size)
 
 
 def test_bvh_query_ray(test, device):
-    test_bvh(test, "ray", device)
+    for leaf_size in [1, 2, 4]:
+        test_bvh(test, "ray", device, leaf_size)
 
 
-def test_gh_288(test, device):
-    num_bounds = 1
-    lowers = ((0.5, -1.0, -1.0),) * num_bounds
-    uppers = ((1.0, 1.0, 1.0),) * num_bounds
+def test_bvh_ray_query_inside_and_outside_bounds(test, device):
+    """Regression test for issue #288: BVH ray queries should detect intersections
+    regardless of whether the ray origin is inside or outside the bounding volumes.
+
+    Previously, rays starting outside the bounds would fail to detect intersections.
+    """
+    # Create a single AABB spanning x=[0.5, 1.0], extending across y and z axes
+    lowers = ((0.5, -1.0, -1.0),)
+    uppers = ((1.0, 1.0, 1.0),)
 
     device_lowers = wp.array(lowers, dtype=wp.vec3f, device=device)
     device_uppers = wp.array(uppers, dtype=wp.vec3f, device=device)
 
     bvh = wp.Bvh(device_lowers, device_uppers)
 
-    bounds_intersected = wp.zeros(shape=num_bounds, dtype=int, device=device)
+    bounds_intersected = wp.zeros(shape=1, dtype=int, device=device)
 
+    # Test both ray origins: outside (x=0.0) and inside (x=0.75) the AABB
     for x in (0.0, 0.75):
         query_start = wp.vec3(x, 0.0, 0.0)
-        query_dir = wp.vec3(1.0, 0.0, 0.0)
+        query_dir = wp.vec3(1.0, 0.0, 0.0)  # Ray pointing in +x direction
 
-        wp.launch(
-            kernel=bvh_query_ray, dim=1, inputs=[bvh.id, query_start, query_dir, bounds_intersected], device=device
-        )
+        wp.launch(bvh_query_ray, dim=1, inputs=[bvh.id, query_start, query_dir, bounds_intersected], device=device)
 
         device_intersected = bounds_intersected.numpy()
-        test.assertEqual(device_intersected.sum(), num_bounds)
+        # Both cases should detect the single intersection
+        test.assertEqual(device_intersected.sum(), 1)
+
+
+def get_random_aabbs(n, center, relative_shift, relative_size, rng):
+    centers = rng.uniform(-0.5, 0.5, size=n * 3).reshape(n, 3) * relative_shift + center
+    diffs = 0.5 * rng.random(n * 3).reshape(n, 3) * relative_size
+
+    lowers = centers - diffs
+    uppers = centers + diffs
+
+    return lowers, uppers
+
+
+@wp.kernel
+def compute_num_contact_with_checksums(
+    lowers: wp.array(dtype=wp.vec3),
+    uppers: wp.array(dtype=wp.vec3),
+    bvh_id: wp.uint64,
+    counts: wp.array(dtype=int),
+    check_sums: wp.array(dtype=int),
+):
+    tid = wp.tid()
+
+    upper = uppers[tid]
+    lower = lowers[tid]
+
+    query = wp.bvh_query_aabb(bvh_id, lower, upper)
+    count = int(0)
+
+    check_sum = int(0)
+    index = int(0)
+    while wp.bvh_query_next(query, index):
+        check_sum = check_sum ^ index
+        count += 1
+
+    counts[tid] = count
+    check_sums[tid] = check_sum
+
+
+def test_capture_bvh_rebuild(test, device):
+    with wp.ScopedDevice(device):
+        rng = np.random.default_rng(123)
+
+        num_item_bounds = 100000
+        item_bound_size = 0.01
+
+        relative_shift = 2
+
+        num_test_bounds = 10000
+        test_bound_relative_size = 0.05
+
+        center = np.array([0.0, 0.0, 0.0])
+
+        item_lowers_np, item_uppers_np = get_random_aabbs(num_item_bounds, center, relative_shift, item_bound_size, rng)
+        item_lowers = wp.array(item_lowers_np, dtype=wp.vec3)
+        item_uppers = wp.array(item_uppers_np, dtype=wp.vec3)
+        bvh_1 = wp.Bvh(item_lowers, item_uppers)
+        item_lowers_2 = wp.zeros_like(item_lowers)
+        item_uppers_2 = wp.zeros_like(item_lowers)
+
+        test_lowers_np, test_uppers_np = get_random_aabbs(
+            num_test_bounds, center, relative_shift, test_bound_relative_size, rng
+        )
+        test_lowers = wp.array(test_lowers_np, dtype=wp.vec3)
+        test_uppers = wp.array(test_uppers_np, dtype=wp.vec3)
+
+        item_lowers_2_np, item_uppers_2_np = get_random_aabbs(
+            num_item_bounds,
+            center,
+            relative_shift,
+            item_bound_size,
+            rng,
+        )
+        item_lowers_2.assign(item_lowers_2_np)
+        item_uppers_2.assign(item_uppers_2_np)
+
+        counts_1 = wp.empty(n=num_test_bounds, dtype=int)
+        checksums_1 = wp.empty(n=num_test_bounds, dtype=int)
+        counts_2 = wp.empty(n=num_test_bounds, dtype=int)
+        checksums_2 = wp.empty(n=num_test_bounds, dtype=int)
+
+        wp.load_module(device=device)
+        with wp.ScopedCapture(force_module_load=False) as capture:
+            wp.copy(item_lowers, item_lowers_2)
+            wp.copy(item_uppers, item_uppers_2)
+            bvh_1.rebuild()
+            wp.launch(
+                compute_num_contact_with_checksums,
+                dim=num_test_bounds,
+                inputs=[test_lowers, test_uppers, bvh_1.id],
+                outputs=[counts_1, checksums_1],
+            )
+
+        cuda_graph = capture.graph
+
+        for _ in range(10):
+            item_lowers_2_np, item_uppers_2_np = get_random_aabbs(
+                num_item_bounds,
+                center,
+                relative_shift,
+                item_bound_size,
+                rng,
+            )
+            item_lowers_2.assign(item_lowers_2_np)
+            item_uppers_2.assign(item_uppers_2_np)
+
+            wp.capture_launch(cuda_graph)
+
+            bvh_2 = wp.Bvh(item_lowers_2, item_uppers_2)
+            wp.launch(
+                compute_num_contact_with_checksums,
+                dim=num_test_bounds,
+                inputs=[test_lowers, test_uppers, bvh_2.id],
+                outputs=[counts_2, checksums_2],
+                device=device,
+            )
+
+            assert_array_equal(counts_1, counts_2)
+            assert_array_equal(checksums_1, checksums_2)
 
 
 devices = get_test_devices()
+cuda_devices = get_cuda_test_devices()
 
 
 class TestBvh(unittest.TestCase):
@@ -193,7 +321,14 @@ class TestBvh(unittest.TestCase):
 
 add_function_test(TestBvh, "test_bvh_aabb", test_bvh_query_aabb, devices=devices)
 add_function_test(TestBvh, "test_bvh_ray", test_bvh_query_ray, devices=devices)
-add_function_test(TestBvh, "test_gh_288", test_gh_288, devices=devices)
+add_function_test(
+    TestBvh,
+    "test_bvh_ray_query_inside_and_outside_bounds",
+    test_bvh_ray_query_inside_and_outside_bounds,
+    devices=devices,
+)
+
+add_function_test(TestBvh, "test_capture_bvh_rebuild", test_capture_bvh_rebuild, devices=cuda_devices)
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/geometry/test_hash_grid.py

@@ -191,7 +191,44 @@ def test_hashgrid_inputs(test, device):
             assert_array_equal(counts_ndim, counts_ref)
 
 
+def test_hashgrid_multiple_streams(test, device):
+    with wp.ScopedDevice(device):
+        points = particle_grid(16, 32, 16, (0.0, 0.3, 0.0), cell_radius * 0.25, 0.1)
+        points_ref = wp.array(points, dtype=wp.vec3)
+        counts_ref = wp.zeros(len(points), dtype=int)
+
+        grid_dim = 64
+        grid_ref = wp.HashGrid(grid_dim, grid_dim, grid_dim)
+        grid_ref.build(points_ref, cell_radius)
+
+        # get reference counts
+        wp.launch(kernel=count_neighbors, dim=len(points), inputs=[grid_ref.id, query_radius, points_ref, counts_ref])
+
+        # create multiple streams
+        num_streams = 10
+        streams = [wp.Stream(device=device) for _ in range(num_streams)]
+        counts_per_stream = [wp.zeros(len(points), dtype=int) for _ in range(num_streams)]
+
+        # test whether HashGrid and radix sort work with multiple streams without race conditions
+        for i in range(num_streams):
+            with wp.ScopedStream(streams[i]):
+                grid = wp.HashGrid(grid_dim, grid_dim, grid_dim)
+                grid.build(points_ref, cell_radius)
+
+                # get counts for this stream
+                wp.launch(
+                    kernel=count_neighbors,
+                    dim=len(points),
+                    inputs=[grid.id, query_radius, points_ref, counts_per_stream[i]],
+                )
+
+        # run this loop after all streams are scheduled to ensure asynchronous behaviour above
+        for i in range(num_streams):
+            assert_array_equal(counts_per_stream[i], counts_ref)
+
+
 devices = get_test_devices()
+cuda_devices = get_cuda_test_devices()
 
 
 class TestHashGrid(unittest.TestCase):
@@ -214,6 +251,7 @@ class TestHashGrid(unittest.TestCase):
 
 add_function_test(TestHashGrid, "test_hashgrid_query", test_hashgrid_query, devices=devices)
 add_function_test(TestHashGrid, "test_hashgrid_inputs", test_hashgrid_inputs, devices=devices)
+add_function_test(TestHashGrid, "test_hashgrid_multiple_streams", test_hashgrid_multiple_streams, devices=cuda_devices)
 
 
 if __name__ == "__main__":

warp/tests/geometry/test_marching_cubes.py

@@ -232,7 +232,6 @@ def test_marching_cubes_differentiable(test, device):
     the sphere's radius.
     """
     node_dim = 64
-    cell_dim = node_dim - 1
     bounds_low = wp.vec3(-1.0, -1.0, -1.0)
     bounds_high = wp.vec3(1.0, 1.0, 1.0)
 

warp/tests/geometry/test_mesh.py

@@ -12,7 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import itertools
 import unittest
 
 import numpy as np
@@ -204,11 +204,13 @@ def test_mesh_query_ray(test, device):
     else:
         constructors = ["sah", "median", "lbvh"]
 
-    for constructor in constructors:
+    leaf_sizes = [1, 2, 4]
+
+    for leaf_size, constructor in itertools.product(leaf_sizes, constructors):
         points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
 
         indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
-        mesh = wp.Mesh(points=points, indices=indices, bvh_constructor=constructor)
+        mesh = wp.Mesh(points=points, indices=indices, bvh_constructor=constructor, bvh_leaf_size=leaf_size)
         expected_sign = -1.0
         wp.launch(
             query_ray_kernel,

warp/tests/geometry/test_mesh_query_aabb.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import itertools
 import os
 import unittest
 
@@ -60,11 +61,6 @@ def compute_bounds(
 def compute_num_contacts(
     lowers: wp.array(dtype=wp.vec3), uppers: wp.array(dtype=wp.vec3), mesh_id: wp.uint64, counts: wp.array(dtype=int)
 ):
-    face_index = int(0)
-    face_u = float(0.0)
-    face_v = float(0.0)
-    sign = float(0.0)
-
     tid = wp.tid()
 
     upper = uppers[tid]
@@ -210,11 +206,6 @@ def compute_num_contact_with_checksums(
     counts: wp.array(dtype=int),
     check_sums: wp.array(dtype=int),
 ):
-    face_index = int(0)
-    face_u = float(0.0)
-    face_v = float(0.0)
-    sign = float(0.0)
-
     tid = wp.tid()
 
     upper = uppers[tid]
@@ -304,12 +295,14 @@ def test_mesh_query_aabb_count_overlap_with_checksum(test, device):
     else:
         constructors = ["sah", "median", "lbvh"]
 
+    leaf_sizes = [1, 2, 4]
+
     points, indices = load_mesh()
     points_wp = wp.array(points, dtype=wp.vec3, device=device)
     indices_wp = wp.array(indices, dtype=int, device=device)
 
-    for constructor in constructors:
-        m = wp.Mesh(points=points_wp, indices=indices_wp, bvh_constructor=constructor)
+    for leaf_size, constructor in itertools.product(leaf_sizes, constructors):
+        m = wp.Mesh(points=points_wp, indices=indices_wp, bvh_constructor=constructor, bvh_leaf_size=leaf_size)
 
         num_test_bounds = 10000
         test_bound_relative_size = 0.01

warp/tests/geometry/test_mesh_query_point.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import itertools
 import math
 import unittest
 
@@ -314,7 +315,9 @@ def test_mesh_query_point(test, device):
     else:
         constructors = ["sah", "median", "lbvh"]
 
-    for constructor in constructors:
+    leaf_sizes = [1, 2, 4]
+
+    for leaf_size, constructor in itertools.product(leaf_sizes, constructors):
         # create mesh
         mesh = wp.Mesh(
             points=mesh_points,
@@ -322,6 +325,7 @@ def test_mesh_query_point(test, device):
             indices=mesh_indices,
             support_winding_number=True,
             bvh_constructor=constructor,
+            bvh_leaf_size=leaf_size,
         )
 
         p = particle_grid(32, 32, 32, np.array([-1.1, -1.1, -1.1]), 0.05, 0.0)
@@ -791,7 +795,6 @@ def point_query_aabb_and_closest(
 
     tri_index = wp.int32(0)
     num_collisions = wp.int32(0)
-    min_dis_to_tris = query_radius
     while wp.mesh_query_aabb_next(query, tri_index):
         t1 = tri_indices[tri_index, 0]
         t2 = tri_indices[tri_index, 1]

warp/tests/geometry/test_mesh_query_ray.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import itertools
 import unittest
 
 import numpy as np
@@ -104,12 +105,20 @@ def test_mesh_query_ray_grad(test, device):
     else:
         constructors = ["sah", "median", "lbvh"]
 
-    for constructor in constructors:
+    leaf_sizes = [1, 2, 4]
+
+    for leaf_size, constructor in itertools.product(leaf_sizes, constructors):
         p = wp.vec3(50.0, 50.0, 0.0)
         D = wp.vec3(0.0, -1.0, 0.0)
 
         # create mesh
-        mesh = wp.Mesh(points=mesh_points, velocities=None, indices=mesh_indices, bvh_constructor=constructor)
+        mesh = wp.Mesh(
+            points=mesh_points,
+            velocities=None,
+            indices=mesh_indices,
+            bvh_constructor=constructor,
+            bvh_leaf_size=leaf_size,
+        )
 
         tape = wp.Tape()
 
@@ -248,6 +257,8 @@ def test_mesh_query_ray_edge(test, device):
     else:
         constructors = ["sah", "median", "lbvh"]
 
+    leaf_sizes = [1, 2, 4]
+
     # Create raycast starts and directions
     xx, yy = np.meshgrid(np.arange(0.1, 0.4, 0.01), np.arange(0.1, 0.4, 0.01))
     xx = xx.flatten().reshape(-1, 1)
@@ -268,11 +279,12 @@ def test_mesh_query_ray_edge(test, device):
 
     triangles = np.array([[1, 0, 2], [1, 2, 3]], dtype=np.int32)
 
-    for constructor in constructors:
+    for leaf_size, constructor in itertools.product(leaf_sizes, constructors):
         mesh = wp.Mesh(
             points=wp.array(vertices, dtype=wp.vec3, device=device),
             indices=wp.array(triangles.flatten(), dtype=int, device=device),
             bvh_constructor=constructor,
+            bvh_leaf_size=leaf_size,
         )
 
         counts = wp.zeros(1, dtype=int, device=device)

warp/tests/geometry/test_volume_write.py

@@ -172,8 +172,8 @@ def test_volume_allocation(test, device):
         device=device,
     )
 
-    assert wp.types.types_equal(volume_a.dtype, wp.float32)
-    assert wp.types.types_equal(volume_b.dtype, wp.float32)
+    assert wp._src.types.types_equal(volume_a.dtype, wp.float32)
+    assert wp._src.types.types_equal(volume_b.dtype, wp.float32)
 
     points = wp.array(points_ref, dtype=wp.vec3, device=device)
     values_a = wp.empty(num_points, dtype=wp.float32, device=device)
@@ -211,8 +211,8 @@ def test_volume_allocate_by_tiles_f(test, device):
     volume_a = wp.Volume.allocate_by_tiles(points_is_d, voxel_size, background_value, translation, device=device)
     volume_b = wp.Volume.allocate_by_tiles(points_ws_d, voxel_size, background_value, translation, device=device)
 
-    assert wp.types.types_equal(volume_a.dtype, wp.float32)
-    assert wp.types.types_equal(volume_b.dtype, wp.float32)
+    assert wp._src.types.types_equal(volume_a.dtype, wp.float32)
+    assert wp._src.types.types_equal(volume_b.dtype, wp.float32)
 
     values_a = wp.empty(num_tiles * 512, dtype=wp.float32, device=device)
     values_b = wp.empty(num_tiles * 512, dtype=wp.float32, device=device)
@@ -243,7 +243,7 @@ def test_volume_allocate_by_tiles_v(test, device):
     points_d = wp.array(points_is, dtype=wp.int32, device=device)
     volume = wp.Volume.allocate_by_tiles(points_d, 0.1, wp.vec3(1, 2, 3), device=device)
 
-    assert wp.types.types_equal(volume.dtype, wp.vec3)
+    assert wp._src.types.types_equal(volume.dtype, wp.vec3)
 
     values = wp.empty(len(points_d) * 512, dtype=wp.vec3, device=device)