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

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:

warp/tests/test_intersect.py

@@ -23,11 +23,28 @@ from warp.tests.unittest_utils import *
 
 @wp.kernel
 def intersect_tri(
-    v0: wp.vec3, v1: wp.vec3, v2: wp.vec3, u0: wp.vec3, u1: wp.vec3, u2: wp.vec3, result: wp.array(dtype=int)
+    v0: wp.vec3,
+    v1: wp.vec3,
+    v2: wp.vec3,
+    u0: wp.vec3,
+    u1: wp.vec3,
+    u2: wp.vec3,
+    result_single_precision: wp.array(dtype=int),
+    result_double_precision: wp.array(dtype=int),
 ):
     tid = wp.tid()
 
-    result[0] = wp.intersect_tri_tri(v0, v1, v2, u0, u1, u2)
+    result_single_precision[0] = wp.intersect_tri_tri(v0, v1, v2, u0, u1, u2)
+
+    u0_d = wp.vec3d(u0)
+    u1_d = wp.vec3d(u1)
+    u2_d = wp.vec3d(u2)
+
+    v0_d = wp.vec3d(v0)
+    v1_d = wp.vec3d(v1)
+    v2_d = wp.vec3d(v2)
+
+    result_double_precision[0] = wp.intersect_tri_tri(v0_d, v1_d, v2_d, u0_d, u1_d, u2_d)
 
 
 def test_intersect_tri(test, device):
@@ -49,13 +66,26 @@ def test_intersect_tri(test, device):
         wp.vec3(-0.5, 0.5, 0.0),
     ]
 
-    result = wp.zeros(1, dtype=int, device=device)
-
-    wp.launch(intersect_tri, dim=1, inputs=[*points_intersect, result], device=device)
-    assert_np_equal(result.numpy(), np.array([1]))
-
-    wp.launch(intersect_tri, dim=1, inputs=[*points_separated, result], device=device)
-    assert_np_equal(result.numpy(), np.array([0]))
+    result_single_precision = wp.zeros(1, dtype=int, device=device)
+    result_double_precision = wp.zeros(1, dtype=int, device=device)
+
+    wp.launch(
+        intersect_tri,
+        dim=1,
+        inputs=[*points_intersect, result_single_precision, result_double_precision],
+        device=device,
+    )
+    assert_np_equal(result_single_precision.numpy(), np.array([1]))
+    assert_np_equal(result_double_precision.numpy(), np.array([1]))
+
+    wp.launch(
+        intersect_tri,
+        dim=1,
+        inputs=[*points_separated, result_single_precision, result_double_precision],
+        device=device,
+    )
+    assert_np_equal(result_single_precision.numpy(), np.array([0]))
+    assert_np_equal(result_double_precision.numpy(), np.array([0]))
 
 
 devices = get_test_devices()

warp/tests/test_large.py

@@ -93,7 +93,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 wptype in wp.types.scalar_types:
+        for wptype in wp._src.types.scalar_types:
             a1 = wp.zeros(shape_tuple, dtype=wptype, device=device)
             assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
 

warp/tests/test_lerp.py

@@ -119,7 +119,9 @@ def test_lerp(test, device):
                 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_size(data_type), dtype=data_type, device=device, requires_grad=True)
+                out = wp.array(
+                    [0] * wp._src.types.type_size(data_type), dtype=data_type, device=device, requires_grad=True
+                )
 
                 with wp.Tape() as tape:
                     wp.launch(kernel, dim=1, inputs=[a, b, t, out], device=device)

warp/tests/test_linear_solvers.py

@@ -18,7 +18,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.optim.linear import bicgstab, cg, cr, gmres, preconditioner
 from warp.tests.unittest_utils import *
 

warp/tests/test_map.py

@@ -18,7 +18,7 @@ import unittest
 import numpy as np
 
 import warp as wp
-import warp.context
+import warp._src.context
 import warp.tests.aux_test_name_clash1 as name_clash_module_1
 import warp.tests.aux_test_name_clash2 as name_clash_module_2
 from warp.tests.unittest_utils import add_function_test, assert_np_equal, get_cuda_test_devices, get_test_devices
@@ -101,6 +101,37 @@ def test_lambda(test, device):
     expected = np.array(np.arange(10) * local_var, dtype=np.float32)
     assert_np_equal(out2.numpy(), expected)
 
+    # inline variable construction which uses parentheses
+    out = wp.map(lambda a: wp.length(wp.vec3(a, a + 1.0, a + 2.0)), a1)
+    expected = np.array([np.sqrt(i * i + (i + 1) * (i + 1) + (i + 2) * (i + 2)) for i in range(10)], dtype=np.float32)
+    assert_np_equal(out.numpy(), expected)
+
+    # multi-line lambda
+    # fmt: off
+    out = wp.map(lambda a: (\
+        a + 1.0 + 2.0 + 3.0 + 4.0 + 5.0 \
+            + 6.0 + 7.0 + 8.0 + 9.0 + 10.0\
+    ), a1)
+    # fmt: on
+    expected = np.array([np.sum(np.arange(1, 11)) + i for i in range(10)], dtype=np.float32)
+    assert_np_equal(out.numpy(), expected)
+
+    # complicated expression with parentheses and line continuation
+    # fmt: off
+    out = wp.map(lambda a: (
+        a + 1.0 + ((
+            + 6.0 + 7.0 + 8.0) + 9.0 + (10.0
+            # an inline comment to make sure it is ignored
+            + 0.0) \
+            + wp.clamp(a, -0.5, 0.5) \
+            + (wp.sin(a) + wp.cos(a))
+    )), a1)
+    # fmt: on
+    expected = np.array(
+        [41.0 + i + np.clip(i, -0.5, 0.5) + (np.sin(i) + np.cos(i)) for i in range(10)], dtype=np.float32
+    )
+    assert_np_equal(out.numpy(), expected, tol=1e-6)
+
 
 def test_multiple_return_values(test, device):
     @wp.func
@@ -152,7 +183,7 @@ def test_custom_struct_operator(test, device):
 
 
 def test_name_clash(test, device):
-    vec5 = wp.types.vector(5, dtype=wp.float32)
+    vec5 = wp._src.types.vector(5, dtype=wp.float32)
 
     @wp.func
     def name_clash_structs_args_func(
@@ -415,8 +446,8 @@ def test_kernel_creation(test, device):
 
 
 def test_graph_capture(test, device):
-    assert warp.context.runtime.driver_version is not None
-    if warp.context.runtime.driver_version < (12, 3):
+    assert warp._src.context.runtime.driver_version is not None
+    if warp._src.context.runtime.driver_version < (12, 3):
         test.skipTest("Module loading during CUDA graph capture is not supported on driver versions < 12.3")
     a_np = np.arange(10, dtype=np.float32)
     b_np = np.arange(1, 11, dtype=np.float32)
@@ -476,6 +507,20 @@ add_function_test(TestMap, "test_graph_capture", test_graph_capture, devices=cud
 add_function_test(TestMap, "test_renamed_warp_module", test_renamed_warp_module, devices=devices)
 
 
+class TestMapDebug(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls._saved_mode = wp.config.mode
+        wp.config.mode = "debug"
+
+    @classmethod
+    def tearDownClass(cls):
+        wp.config.mode = cls._saved_mode
+
+
+add_function_test(TestMapDebug, "test_mixed_inputs", test_mixed_inputs, devices=devices)
+add_function_test(TestMapDebug, "test_kernel_creation", test_kernel_creation, devices=devices)
+
 if __name__ == "__main__":
     wp.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_mat.py

@@ -67,17 +67,17 @@ def test_shape_mismatch(test, device):
 
 
 def test_py_arithmetic_ops(test, device, dtype):
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
 
     def make_mat(*args):
-        if wptype in wp.types.int_types:
+        if wptype in wp._src.types.int_types:
             # Cast to the correct integer type to simulate wrapping.
             return tuple(tuple(wptype._type_(x).value for x in row) for row in args)
 
         return args
 
     def make_vec(*args):
-        if wptype in wp.types.int_types:
+        if wptype in wp._src.types.int_types:
             # Cast to the correct integer type to simulate wrapping.
             return tuple(wptype._type_(x).value for x in args)
 
@@ -116,11 +116,11 @@ def test_negation(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    mat22 = wp._src.types.matrix(shape=(2, 2), dtype=wptype)
+    mat33 = wp._src.types.matrix(shape=(3, 3), dtype=wptype)
+    mat44 = wp._src.types.matrix(shape=(4, 4), dtype=wptype)
+    mat55 = wp._src.types.matrix(shape=(5, 5), dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -205,12 +205,12 @@ def test_matmul(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-12,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat23 = wp.types.matrix(shape=(2, 3), dtype=wptype)
-    mat32 = wp.types.matrix(shape=(3, 2), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    mat22 = wp._src.types.matrix(shape=(2, 2), dtype=wptype)
+    mat33 = wp._src.types.matrix(shape=(3, 3), dtype=wptype)
+    mat23 = wp._src.types.matrix(shape=(2, 3), dtype=wptype)
+    mat32 = wp._src.types.matrix(shape=(3, 2), dtype=wptype)
+    mat44 = wp._src.types.matrix(shape=(4, 4), dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -276,11 +276,11 @@ def test_subtraction(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    mat22 = wp._src.types.matrix(shape=(2, 2), dtype=wptype)
+    mat33 = wp._src.types.matrix(shape=(3, 3), dtype=wptype)
+    mat44 = wp._src.types.matrix(shape=(4, 4), dtype=wptype)
+    mat55 = wp._src.types.matrix(shape=(5, 5), dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -397,10 +397,10 @@ def test_determinant(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    mat22 = wp._src.types.matrix(shape=(2, 2), dtype=wptype)
+    mat33 = wp._src.types.matrix(shape=(3, 3), dtype=wptype)
+    mat44 = wp._src.types.matrix(shape=(4, 4), dtype=wptype)
 
     def check_mat_det(
         v2: wp.array(dtype=mat22),
@@ -535,8 +535,8 @@ def test_determinant(test, device, dtype, register_kernels=False):
 #         np.float64: 1.0e-8,
 #     }.get(dtype, 0)
 #
-#     wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-#     mat55 = wp.types.vector(shape=(5, 5), dtype=wptype)
+#     wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+#     mat55 = wp._src.types.vector(shape=(5, 5), dtype=wptype)
 #
 #     output_select_kernel = get_select_kernel(wptype)
 #
@@ -590,10 +590,10 @@ def test_inverse(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    mat22 = wp._src.types.matrix(shape=(2, 2), dtype=wptype)
+    mat33 = wp._src.types.matrix(shape=(3, 3), dtype=wptype)
+    mat44 = wp._src.types.matrix(shape=(4, 4), dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -760,9 +760,9 @@ def test_svd(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-12,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    vec3 = wp.types.vector(length=3, dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
+    mat33 = wp._src.types.matrix(shape=(3, 3), dtype=wptype)
 
     def check_mat_svd(
         m3: wp.array(dtype=mat33),
@@ -879,9 +879,9 @@ def test_svd_2D(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-12,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    vec2 = wp.types.vector(length=2, dtype=wptype)
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    vec2 = wp._src.types.vector(length=2, dtype=wptype)
+    mat22 = wp._src.types.matrix(shape=(2, 2), dtype=wptype)
 
     def check_mat_svd2(
         m2: wp.array(dtype=mat22),
@@ -1027,8 +1027,8 @@ def test_qr(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-12,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    mat33 = wp._src.types.matrix(shape=(3, 3), dtype=wptype)
 
     def check_mat_qr(
         m3: wp.array(dtype=mat33),
@@ -1139,9 +1139,9 @@ def test_eig(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-5,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    vec3 = wp.types.vector(length=3, dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
+    mat33 = wp._src.types.matrix(shape=(3, 3), dtype=wptype)
 
     def check_mat_eig(
         m3: wp.array(dtype=mat33),
@@ -1251,8 +1251,8 @@ def test_skew(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    vec3 = wp.types.vector(length=3, dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -1283,7 +1283,6 @@ def test_skew(test, device, dtype, register_kernels=False):
 
     # make sure it gives you a cross product matrix:
     crossprodmat = outcomponents.numpy().reshape(3, 3)
-    v = np.array([1, 0, 0])
     assert_np_equal(
         np.matmul(crossprodmat, np.array([1, 0, 0])).reshape(-1),
         2 * np.cross(v3.numpy()[0], np.array([1, 0, 0])),
@@ -1353,9 +1352,9 @@ def test_transform_point(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    vec3 = wp.types.vector(length=3, dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
+    mat44 = wp._src.types.matrix(shape=(4, 4), dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -1414,9 +1413,9 @@ def test_transform_vector(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    vec3 = wp.types.vector(length=3, dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
+    mat44 = wp._src.types.matrix(shape=(4, 4), dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -1466,22 +1465,13 @@ def test_transform_vector(test, device, dtype, register_kernels=False):
 
 
 @wp.kernel
-def test_matrix_mutation(expected: wp.types.matrix(shape=(10, 3), dtype=float)):
+def test_matrix_mutation(expected: wp._src.types.matrix(shape=(10, 3), dtype=float)):
     m = wp.matrix(shape=(10, 3), dtype=float)
 
     # test direct element indexing
     m[0, 0] = 1.0
-    m[0, 1] = 2.0
-    m[0, 2] = 3.0
-
-    # The nested indexing (matrix->vector->scalar) below does not
-    # currently modify m because m[0] returns row vector by
-    # value rather than reference, this is different from NumPy
-    # which always returns by ref. Not clear how we can support
-    # this as well as auto-diff.
-
-    # m[0][1] = 2.0
-    # m[0][2] = 3.0
+    m[0][1] = 2.0
+    m[0][2] = 3.0
 
     # test setting rows
     for i in range(1, 10):
@@ -3443,7 +3433,7 @@ class TestMat(unittest.TestCase):
         self.assertSequenceEqual(m, ((0.0, 1.0), (2.0, 3.0)))
 
 
-mat103 = wp.types.matrix(shape=(10, 3), dtype=float)
+mat103 = wp._src.types.matrix(shape=(10, 3), dtype=float)
 add_kernel_test(
     TestMat,
     test_matrix_mutation,

warp/tests/test_mat_constructors.py

@@ -419,11 +419,10 @@ def test_quat_constructor(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    vec4 = wp.types.vector(length=4, dtype=wptype)
-    vec3 = wp.types.vector(length=3, dtype=wptype)
-    quat = wp.types.quaternion(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    vec4 = wp._src.types.vector(length=4, dtype=wptype)
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
+    quat = wp._src.types.quaternion(dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 

warp/tests/test_mat_lite.py

@@ -81,7 +81,7 @@ def test_constructors_default_precision():
 
 
 @wp.kernel
-def test_matrix_mutation(expected: wp.types.matrix(shape=(10, 3), dtype=float)):
+def test_matrix_mutation(expected: wp._src.types.matrix(shape=(10, 3), dtype=float)):
     m = wp.matrix(shape=(10, 3), dtype=float)
 
     # test direct element indexing