warp-lang 1.9.1__py3-none-manylinux_2_34_aarch64.whl → 1.10.0rc2__py3-none-manylinux_2_34_aarch64.whl

warp/tests/test_examples.py

@@ -44,15 +44,13 @@ from typing import Any, Dict, Optional, Type
 
 import warp as wp
 import warp.tests.unittest_utils
+from warp._src.utils import check_p2p
 from warp.tests.unittest_utils import (
     USD_AVAILABLE,
     get_selected_cuda_test_devices,
     get_test_devices,
     sanitize_identifier,
 )
-from warp.utils import check_p2p
-
-wp.init()  # For wp.context.runtime.core.wp_is_debug_enabled()
 
 
 def _build_command_line_options(test_options: Dict[str, Any]) -> list:
@@ -280,25 +278,6 @@ class TestOptimExamples(unittest.TestCase):
     pass
 
 
-add_example_test(
-    TestOptimExamples,
-    name="optim.example_bounce",
-    devices=test_devices,
-    test_options_cpu={"train_iters": 3},
-)
-add_example_test(
-    TestOptimExamples,
-    name="optim.example_drone",
-    devices=test_devices,
-    test_options={"headless": True},
-    test_options_cpu={"num_frames": 10},
-)
-add_example_test(
-    TestOptimExamples,
-    name="optim.example_cloth_throw",
-    devices=test_devices,
-    test_options_cpu={"train_iters": 3},
-)
 add_example_test(
     TestOptimExamples,
     name="optim.example_diffray",
@@ -312,68 +291,6 @@ add_example_test(
     devices=cuda_test_devices,
     test_options={"headless": True, "train_iters": 5, "num_frames": 300, "pillow_required": True},
 )
-add_example_test(TestOptimExamples, name="optim.example_inverse_kinematics", devices=test_devices)
-add_example_test(
-    TestOptimExamples,
-    name="optim.example_inverse_kinematics_torch",
-    devices=test_devices,
-    test_options={"torch_required": True},
-)
-add_example_test(TestOptimExamples, name="optim.example_spring_cage", devices=test_devices)
-add_example_test(
-    TestOptimExamples,
-    name="optim.example_trajectory",
-    devices=test_devices,
-    test_options={"headless": True, "train_iters": 50},
-)
-add_example_test(
-    TestOptimExamples,
-    name="optim.example_softbody_properties",
-    devices=test_devices,
-    test_options_cuda={
-        "train_iters": 1 if warp.context.runtime.core.wp_is_debug_enabled() else 3,
-    },
-    test_options_cpu={"train_iters": 1},
-)
-
-
-class TestSimExamples(unittest.TestCase):
-    pass
-
-
-add_example_test(TestSimExamples, name="sim.example_cartpole", devices=test_devices)
-add_example_test(
-    TestSimExamples,
-    name="sim.example_cloth",
-    devices=test_devices,
-    test_options={"usd_required": True},
-    test_options_cpu={"num_frames": 10},
-)
-add_example_test(
-    TestSimExamples, name="sim.example_granular", devices=test_devices, test_options_cpu={"num_frames": 10}
-)
-add_example_test(TestSimExamples, name="sim.example_granular_collision_sdf", devices=cuda_test_devices)
-add_example_test(TestSimExamples, name="sim.example_jacobian_ik", devices=test_devices)
-add_example_test(TestSimExamples, name="sim.example_particle_chain", devices=test_devices)
-add_example_test(
-    TestSimExamples, name="sim.example_quadruped", devices=test_devices, test_options_cpu={"num_frames": 100}
-)
-add_example_test(TestSimExamples, name="sim.example_rigid_chain", devices=test_devices)
-add_example_test(
-    TestSimExamples,
-    name="sim.example_rigid_contact",
-    devices=test_devices,
-    test_options={"usd_required": True},
-    test_options_cpu={"num_frames": 3},
-)
-add_example_test(
-    TestSimExamples, name="sim.example_rigid_soft_contact", devices=test_devices, test_options_cpu={"num_frames": 10}
-)
-add_example_test(TestSimExamples, name="sim.example_rigid_force", devices=test_devices)
-add_example_test(TestSimExamples, name="sim.example_rigid_gyroscopic", devices=test_devices)
-add_example_test(
-    TestSimExamples, name="sim.example_soft_body", devices=test_devices, test_options_cpu={"num_frames": 10}
-)
 
 
 class TestFemExamples(unittest.TestCase):
@@ -486,6 +403,21 @@ add_example_test(
     test_options={"headless": True, "resolution": 16, "num_steps": 2},
 )
 
+
+class TestTileExamples(unittest.TestCase):
+    pass
+
+
+add_example_test(
+    TestTileExamples,
+    name="tile.example_tile_mcgp",
+    devices=test_devices,
+    test_options={"height": 128, "headless": True, "usd_required": True},
+    test_options_cpu={"slices": 2},
+    test_options_cuda={"slices": 5},
+)
+
+
 if __name__ == "__main__":
     # force rebuild of all kernels
     wp.clear_kernel_cache()

warp/tests/test_fabricarray.py

@@ -26,14 +26,14 @@ from warp.tests.unittest_utils import *
 
 # types to test fabric arrays
 _fabric_types = [
-    *wp.types.scalar_types,
-    *[wp.types.vector(2, T) for T in wp.types.scalar_types],
-    *[wp.types.vector(3, T) for T in wp.types.scalar_types],
-    *[wp.types.vector(4, T) for T in wp.types.scalar_types],
-    *[wp.types.matrix((2, 2), T) for T in wp.types.scalar_types],
-    *[wp.types.matrix((3, 3), T) for T in wp.types.scalar_types],
-    *[wp.types.matrix((4, 4), T) for T in wp.types.scalar_types],
-    *[wp.types.quaternion(T) for T in wp.types.float_types],
+    *wp._src.types.scalar_types,
+    *[wp._src.types.vector(2, T) for T in wp._src.types.scalar_types],
+    *[wp._src.types.vector(3, T) for T in wp._src.types.scalar_types],
+    *[wp._src.types.vector(4, T) for T in wp._src.types.scalar_types],
+    *[wp._src.types.matrix((2, 2), T) for T in wp._src.types.scalar_types],
+    *[wp._src.types.matrix((3, 3), T) for T in wp._src.types.scalar_types],
+    *[wp._src.types.matrix((4, 4), T) for T in wp._src.types.scalar_types],
+    *[wp._src.types.quaternion(T) for T in wp._src.types.float_types],
 ]
 
 
@@ -114,7 +114,7 @@ def _create_fabric_array_interface(data: wp.array, attrib: str, bucket_sizes: li
         # empty data array
         bucket_sizes = []
 
-    dtype_size = wp.types.type_size_in_bytes(data.dtype)
+    dtype_size = wp._src.types.type_size_in_bytes(data.dtype)
     p = int(data.ptr) if data.ptr else 0
     pointers = []
     counts = []
@@ -189,10 +189,10 @@ def _create_fabric_array_array_interface(data: list, attrib: str, bucket_sizes:
         _array_lengths.append(data[i].size)
 
     array_pointers = wp.array(_array_pointers, dtype=wp.uint64, device=device)
-    pointer_size = wp.types.type_size_in_bytes(array_pointers.dtype)
+    pointer_size = wp._src.types.type_size_in_bytes(array_pointers.dtype)
 
     lengths = wp.array(_array_lengths, dtype=wp.uint64, device=device)
-    length_size = wp.types.type_size_in_bytes(lengths.dtype)
+    length_size = wp._src.types.type_size_in_bytes(lengths.dtype)
 
     p_pointers = int(array_pointers.ptr)
     p_lengths = int(lengths.ptr)
@@ -296,9 +296,9 @@ def fa_generic_dtype_kernel_indexed(a: wp.indexedfabricarray(dtype=Any), b: wp.i
 def test_fabricarray_generic_dtype(test, device):
     for T in _fabric_types:
         if hasattr(T, "_wp_scalar_type_"):
-            nptype = wp.types.warp_type_to_np_dtype[T._wp_scalar_type_]
+            nptype = wp._src.types.warp_type_to_np_dtype[T._wp_scalar_type_]
         else:
-            nptype = wp.types.warp_type_to_np_dtype[T]
+            nptype = wp._src.types.warp_type_to_np_dtype[T]
 
         data = wp.array(data=np.arange(10, dtype=nptype), device=device)
         data_iface = _create_fabric_array_interface(data, "foo", copy=True)
@@ -345,9 +345,9 @@ def fa_generic_array_kernel(a: Any, b: Any):
 def test_fabricarray_generic_array(test, device):
     for T in _fabric_types:
         if hasattr(T, "_wp_scalar_type_"):
-            nptype = wp.types.warp_type_to_np_dtype[T._wp_scalar_type_]
+            nptype = wp._src.types.warp_type_to_np_dtype[T._wp_scalar_type_]
         else:
-            nptype = wp.types.warp_type_to_np_dtype[T]
+            nptype = wp._src.types.warp_type_to_np_dtype[T]
 
         data = wp.array(data=np.arange(100, dtype=nptype), device=device)
         data_iface = _create_fabric_array_interface(data, "foo", copy=True)
@@ -393,9 +393,9 @@ def test_fabricarray_empty(test, device):
         # scalar, vector, or matrix
         if ncols > 0:
             if nrows > 0:
-                wptype = wp.types.matrix((nrows, ncols), wptype)
+                wptype = wp._src.types.matrix((nrows, ncols), wptype)
             else:
-                wptype = wp.types.vector(ncols, wptype)
+                wptype = wp._src.types.vector(ncols, wptype)
             dtype_shape = wptype._shape_
         else:
             dtype_shape = ()
@@ -461,7 +461,7 @@ def test_fabricarray_empty(test, device):
         test.assertEqual(ifa.list(), [])
 
     # test with scalars, vectors, and matrices
-    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
+    for nptype, wptype in wp._src.types.np_dtype_to_warp_type.items():
         # scalars
         test_empty_ops(0, 0, wptype, nptype)
 
@@ -473,7 +473,7 @@ def test_fabricarray_empty(test, device):
 
 
 def test_fabricarray_fill_scalar(test, device):
-    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
+    for nptype, wptype in wp._src.types.np_dtype_to_warp_type.items():
         # create a data array
         data = wp.zeros(100, dtype=wptype, device=device)
         iface = _create_fabric_array_interface(data, "foo", copy=True)
@@ -489,7 +489,7 @@ def test_fabricarray_fill_scalar(test, device):
         fa.zero_()
         assert_np_equal(fa.numpy(), np.zeros(fa.shape, dtype=nptype))
 
-        if wptype in wp.types.float_types:
+        if wptype in wp._src.types.float_types:
             # fill with float value
             fill_value = 13.37
             fa.fill_(fill_value)
@@ -524,7 +524,7 @@ def test_fabricarray_fill_scalar(test, device):
         assert_np_equal(ifa.numpy(), np.zeros(ifa.shape, dtype=nptype))
         assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))
 
-        if wptype in wp.types.float_types:
+        if wptype in wp._src.types.float_types:
             # fill with float value
             fill_value = 13.37
             ifa.fill_(fill_value)
@@ -541,13 +541,13 @@ def test_fabricarray_fill_scalar(test, device):
 def test_fabricarray_fill_vector(test, device):
     # test filling a vector array with scalar or vector values (vec_type, list, or numpy array)
 
-    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
+    for nptype, wptype in wp._src.types.np_dtype_to_warp_type.items():
         # vector types
         vector_types = [
-            wp.types.vector(2, wptype),
-            wp.types.vector(3, wptype),
-            wp.types.vector(4, wptype),
-            wp.types.vector(5, wptype),
+            wp._src.types.vector(2, wptype),
+            wp._src.types.vector(3, wptype),
+            wp._src.types.vector(4, wptype),
+            wp._src.types.vector(5, wptype),
         ]
 
         for vec_type in vector_types:
@@ -593,7 +593,7 @@ def test_fabricarray_fill_vector(test, device):
             fa.fill_(fill_vec)
             assert_np_equal(fa.numpy(), expected)
 
-            if wptype in wp.types.float_types:
+            if wptype in wp._src.types.float_types:
                 # fill with float scalar
                 fill_value = 13.37
                 fa.fill_(fill_value)
@@ -661,7 +661,7 @@ def test_fabricarray_fill_vector(test, device):
             assert_np_equal(ifa.numpy(), expected)
             assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
 
-            if wptype in wp.types.float_types:
+            if wptype in wp._src.types.float_types:
                 # fill with float scalar
                 fill_value = 13.37
                 ifa.fill_(fill_value)
@@ -682,14 +682,14 @@ def test_fabricarray_fill_vector(test, device):
 def test_fabricarray_fill_matrix(test, device):
     # test filling a matrix array with scalar or matrix values (mat_type, nested list, or 2d numpy array)
 
-    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
+    for nptype, wptype in wp._src.types.np_dtype_to_warp_type.items():
         # matrix types
         matrix_types = [
             # square matrices only
-            wp.types.matrix((2, 2), wptype),
-            wp.types.matrix((3, 3), wptype),
-            wp.types.matrix((4, 4), wptype),
-            wp.types.matrix((5, 5), wptype),
+            wp._src.types.matrix((2, 2), wptype),
+            wp._src.types.matrix((3, 3), wptype),
+            wp._src.types.matrix((4, 4), wptype),
+            wp._src.types.matrix((5, 5), wptype),
         ]
 
         for mat_type in matrix_types:
@@ -896,9 +896,9 @@ def fa_generic_sums_kernel_indexed(a: wp.indexedfabricarrayarray(dtype=Any), sum
 def test_fabricarrayarray(test, device):
     for T in _fabric_types:
         if hasattr(T, "_wp_scalar_type_"):
-            nptype = wp.types.warp_type_to_np_dtype[T._wp_scalar_type_]
+            nptype = wp._src.types.warp_type_to_np_dtype[T._wp_scalar_type_]
         else:
-            nptype = wp.types.warp_type_to_np_dtype[T]
+            nptype = wp._src.types.warp_type_to_np_dtype[T]
 
         n = 100
 

warp/tests/test_fast_math.py

@@ -21,8 +21,6 @@ from warp.tests.unittest_utils import *
 
 @wp.kernel
 def test_pow(e: float, expected: float):
-    tid = wp.tid()
-
     y = wp.pow(-2.0, e)
 
     # Since equality comparisons with NaN's are false, we have to do something manually

warp/tests/test_fem.py

@@ -22,19 +22,18 @@ import numpy as np
 
 import warp as wp
 import warp.fem as fem
-from warp.fem import Coords, D, Domain, Field, Sample, curl, div, grad, integrand, normal
+from warp._src.fem.geometry.closest_point import project_on_tet_at_origin, project_on_tri_at_origin
+from warp._src.fem.space import shape
+from warp.fem import Coords, D, Domain, Field, Sample, curl, div, grad, integrand, make_free_sample, normal
 from warp.fem.cache import dynamic_kernel
-from warp.fem.geometry.closest_point import project_on_tet_at_origin, project_on_tri_at_origin
 from warp.fem.linalg import inverse_qr, spherical_part, symmetric_eigenvalues_qr, symmetric_part
-from warp.fem.space import shape
-from warp.fem.types import make_free_sample
 from warp.fem.utils import (
     grid_to_hexes,
     grid_to_quads,
     grid_to_tets,
     grid_to_tris,
 )
-from warp.sparse import bsr_zeros
+from warp.sparse import bsr_set_zero, bsr_zeros
 from warp.tests.unittest_utils import *
 
 vec6f = wp.vec(length=6, dtype=float)
@@ -586,7 +585,7 @@ def _launch_test_geometry_kernel(geo: fem.Geometry, device):
         fem.interpolate(
             _test_geo_sides,
             quadrature=side_quadrature,
-            values={"side_measures": side_measures, "ref_measure": geo.reference_side().measure()},
+            values={"side_measures": side_measures, "ref_measure": geo.reference_side().prototype.measure()},
         )
 
         if geo.side_normal is not None:
@@ -775,7 +774,7 @@ def _refinement_field(x: wp.vec3):
 def test_adaptive_nanogrid(test, device):
     # 3 res-1 voxels, 8 res-0 voxels
 
-    if platform.system() == "Windows" or (device.is_cuda and wp.context.runtime.toolkit_version[0] == 11):
+    if platform.system() == "Windows" or (device.is_cuda and wp._src.context.runtime.toolkit_version[0] == 11):
         test.skipTest("Skipping test due to NVRTC bug on CUDA 11 and Windows")
 
     res0 = wp.array(
@@ -1045,7 +1044,7 @@ def test_closest_point_queries(test, device):
 
 
 def test_regular_quadrature(test, device):
-    from warp.fem.geometry.element import LinearEdge, Polynomial, Triangle
+    from warp._src.fem.geometry.element import LinearEdge, Polynomial, Triangle
 
     for family in Polynomial:
         # test integrating monomials
@@ -1573,7 +1572,9 @@ def test_particle_quadratures(test, device):
     domain = fem.Cells(geo)
 
     # Explicit quadrature
-    points, weights = domain.reference_element().instantiate_quadrature(order=4, family=fem.Polynomial.GAUSS_LEGENDRE)
+    points, weights = domain.reference_element().prototype.instantiate_quadrature(
+        order=4, family=fem.Polynomial.GAUSS_LEGENDRE
+    )
     points_per_cell = len(points)
 
     points = points * domain.element_count()
@@ -2046,7 +2047,7 @@ def test_array_axpy(test, device):
 
     tape = wp.Tape()
     with tape:
-        fem.utils.array_axpy(x=x, y=y, alpha=alpha, beta=beta)
+        fem.linalg.array_axpy(x=x, y=y, alpha=alpha, beta=beta)
 
     assert_np_equal(x.numpy(), np.full(N, 2.0))
     assert_np_equal(y.numpy(), alpha * x.numpy() + beta * np.arange(N))
@@ -2097,6 +2098,50 @@ def test_integrate_high_order(test_field, device):
         assert_np_equal(h0.values[:h0_nnz].numpy(), h1.values[:h1_nnz].numpy(), tol=1.0e-6)
 
 
+def test_capturability(test, device):
+    A = bsr_zeros(0, 0, block_type=wp.float32, device=device)
+
+    def test_body():
+        geo = fem.Grid3D(res=(4, 4, 4))
+        space = fem.make_polynomial_space(geo, degree=1)
+
+        cell_mask = wp.zeros(geo.cell_count(), dtype=int, device=device)
+        cell_mask[16:32].fill_(1)
+
+        geo_partition = fem.ExplicitGeometryPartition(geo, cell_mask, max_cell_count=32, max_side_count=0)
+        space_partition = fem.make_space_partition(space, geo_partition, with_halo=False, max_node_count=64)
+
+        test_field = fem.make_test(space, space_partition=space_partition)
+        trial_field = fem.make_trial(space, space_partition=space_partition)
+        bsr_set_zero(
+            A,
+            rows_of_blocks=test_field.space_partition.node_count(),
+            cols_of_blocks=trial_field.space_partition.node_count(),
+        )
+        fem.integrate(
+            bilinear_form,
+            fields={"v": test_field, "u": trial_field},
+            kernel_options={"enable_backward": False},
+            output=A,
+        )
+
+    with wp.ScopedDevice(device):
+        test_body()
+        assert A.shape == (64, 64)
+        nnz_ref = A.nnz_sync()
+        values_ref = A.values.numpy()[:nnz_ref]
+        columns_ref = A.columns.numpy()[:nnz_ref]
+        bsr_set_zero(A)
+        assert A.nnz_sync() == 0
+
+        with wp.ScopedCapture() as capture:
+            test_body()
+        wp.capture_launch(capture.graph)
+        assert A.nnz_sync() == nnz_ref
+        assert_np_equal(A.values.numpy()[:nnz_ref], values_ref)
+        assert_np_equal(A.columns.numpy()[:nnz_ref], columns_ref)
+
+
 devices = get_test_devices()
 cuda_devices = get_selected_cuda_test_devices()
 
@@ -2152,6 +2197,7 @@ add_function_test(TestFemShapeFunctions, "test_cube_shape_functions", test_cube_
 add_function_test(TestFemShapeFunctions, "test_tri_shape_functions", test_tri_shape_functions)
 add_function_test(TestFemShapeFunctions, "test_tet_shape_functions", test_tet_shape_functions)
 
+add_function_test(TestFemShapeFunctions, "test_capturability", test_capturability, devices=cuda_devices)
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/test_fixedarray.py

@@ -115,7 +115,7 @@ def test_error_invalid_func_return_annotation(test, device):
         arr = func()
 
     with test.assertRaisesRegex(
-        wp.codegen.WarpCodegenError,
+        wp._src.codegen.WarpCodegenError,
         r"The function `func` returns a fixed-size array whereas it has its return type annotated as `Array\[int32\]`.$",
     ):
         wp.launch(kernel, 1, device=device)
@@ -152,8 +152,8 @@ def test_capture_if_kernel():
 def test_capture_if(test, device):
     if (
         not wp.get_device(device).is_cuda
-        or wp.context.runtime.toolkit_version < (12, 4)
-        or wp.context.runtime.driver_version < (12, 4)
+        or wp._src.context.runtime.toolkit_version < (12, 4)
+        or wp._src.context.runtime.driver_version < (12, 4)
     ):
         return
 

warp/tests/test_func.py

@@ -396,7 +396,7 @@ class TestFunc(unittest.TestCase):
         t = wp.transform(*wp.transform(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0))
         self.assertSequenceEqual(t, (1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0))
 
-        transformf = wp.types.transformation(dtype=float)
+        transformf = wp._src.types.transformation(dtype=float)
 
         t = wp.transformf((1.0, 2.0, 3.0), (4.0, 5.0, 6.0, 7.0))
         self.assertSequenceEqual(
@@ -408,6 +408,12 @@ class TestFunc(unittest.TestCase):
             (-1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0),
         )
 
+        with self.assertRaisesRegex(
+            TypeError,
+            r"^got an unexpected keyword argument 'pos'$",
+        ):
+            wp.transform(pos=wp.vec3(), rot=wp.quat())
+
         f = wp.sin(math.pi * 0.5)
         self.assertAlmostEqual(f, 1.0, places=3)
 
@@ -443,10 +449,7 @@ class TestFunc(unittest.TestCase):
             a * b
 
     def test_cpython_call_user_function_with_error(self):
-        with self.assertRaisesRegex(
-            ZeroDivisionError,
-            "float division by zero",
-        ):
+        with self.assertRaises(ZeroDivisionError):
             divide_by_zero(1.0)
 
     def test_cpython_call_user_function_with_wrong_argument_types(self):

warp/tests/test_generics.py

@@ -521,7 +521,7 @@ def test_type_attribute_error(test, device):
         a = wp.vec3(0.0)
         _ = a.dtype.shape
 
-    with test.assertRaisesRegex(AttributeError, r"`shape` is not an attribute of '<class 'warp.types.float32'>'"):
+    with test.assertRaisesRegex(AttributeError, r"`shape` is not an attribute of '<class 'warp._src.types.float32'>'"):
         wp.launch(
             kernel,
             dim=1,

warp/tests/test_indexedarray.py

@@ -268,9 +268,9 @@ def test_indexedarray_mixed(test, device):
     wp.launch(kernel_3d, dim=iarr.shape, inputs=[iarr, expected_arr], device=device)
 
 
-vec2i = wp.types.vector(length=2, dtype=wp.int32)
-vec3i = wp.types.vector(length=3, dtype=wp.int32)
-vec4i = wp.types.vector(length=4, dtype=wp.int32)
+vec2i = wp._src.types.vector(length=2, dtype=wp.int32)
+vec3i = wp._src.types.vector(length=3, dtype=wp.int32)
+vec4i = wp._src.types.vector(length=4, dtype=wp.int32)
 
 
 @wp.kernel
@@ -571,13 +571,13 @@ def test_indexedarray_empty(test, device):
         data_shape = (1,) * ndim
         dtype_shape = ()
 
-        if wptype in wp.types.scalar_types:
+        if wptype in wp._src.types.scalar_types:
             # scalar, vector, or matrix
             if ncols > 0:
                 if nrows > 0:
-                    wptype = wp.types.matrix((nrows, ncols), wptype)
+                    wptype = wp._src.types.matrix((nrows, ncols), wptype)
                 else:
-                    wptype = wp.types.vector(ncols, wptype)
+                    wptype = wp._src.types.vector(ncols, wptype)
                 dtype_shape = wptype._shape_
             fill_value = wptype(42)
         else:
@@ -620,7 +620,7 @@ def test_indexedarray_empty(test, device):
 
     for ndim in range(1, 5):
         # test with scalars, vectors, and matrices
-        for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
+        for nptype, wptype in wp._src.types.np_dtype_to_warp_type.items():
             # scalars
             test_empty_ops(ndim, 0, 0, wptype, nptype)
 
@@ -643,7 +643,7 @@ def test_indexedarray_empty(test, device):
 def test_indexedarray_fill_scalar(test, device):
     dim_x = 4
 
-    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
+    for nptype, wptype in wp._src.types.np_dtype_to_warp_type.items():
         data1 = wp.zeros(dim_x, dtype=wptype, device=device)
         data2 = wp.zeros((dim_x, dim_x), dtype=wptype, device=device)
         data3 = wp.zeros((dim_x, dim_x, dim_x), dtype=wptype, device=device)
@@ -684,7 +684,7 @@ def test_indexedarray_fill_scalar(test, device):
         assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
         assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
 
-        if wptype in wp.types.float_types:
+        if wptype in wp._src.types.float_types:
             # fill with float value
             fill_value = 13.37
 
@@ -717,13 +717,13 @@ def test_indexedarray_fill_vector(test, device):
 
     dim_x = 4
 
-    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
+    for nptype, wptype in wp._src.types.np_dtype_to_warp_type.items():
         # vector types
         vector_types = [
-            wp.types.vector(2, wptype),
-            wp.types.vector(3, wptype),
-            wp.types.vector(4, wptype),
-            wp.types.vector(5, wptype),
+            wp._src.types.vector(2, wptype),
+            wp._src.types.vector(3, wptype),
+            wp._src.types.vector(4, wptype),
+            wp._src.types.vector(5, wptype),
         ]
 
         for vec_type in vector_types:
@@ -825,7 +825,7 @@ def test_indexedarray_fill_vector(test, device):
             assert_np_equal(a3.numpy(), expected3)
             assert_np_equal(a4.numpy(), expected4)
 
-            if wptype in wp.types.float_types:
+            if wptype in wp._src.types.float_types:
                 # fill with float scalar
                 fill_value = 13.37
 
@@ -863,19 +863,19 @@ def test_indexedarray_fill_matrix(test, device):
 
     dim_x = 4
 
-    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
+    for nptype, wptype in wp._src.types.np_dtype_to_warp_type.items():
         # matrix types
         matrix_types = [
             # square matrices
-            wp.types.matrix((2, 2), wptype),
-            wp.types.matrix((3, 3), wptype),
-            wp.types.matrix((4, 4), wptype),
-            wp.types.matrix((5, 5), wptype),
+            wp._src.types.matrix((2, 2), wptype),
+            wp._src.types.matrix((3, 3), wptype),
+            wp._src.types.matrix((4, 4), wptype),
+            wp._src.types.matrix((5, 5), wptype),
             # non-square matrices
-            wp.types.matrix((2, 3), wptype),
-            wp.types.matrix((3, 2), wptype),
-            wp.types.matrix((3, 4), wptype),
-            wp.types.matrix((4, 3), wptype),
+            wp._src.types.matrix((2, 3), wptype),
+            wp._src.types.matrix((3, 2), wptype),
+            wp._src.types.matrix((3, 4), wptype),
+            wp._src.types.matrix((4, 3), wptype),
         ]
 
         for mat_type in matrix_types: