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

warp/tests/test_spatial.py

@@ -55,9 +55,9 @@ def test_spatial_vector_constructors(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)
-    spatial_vector = wp.types.vector(length=6, dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
+    spatial_vector = wp._src.types.vector(length=6, dtype=wptype)
 
     def check_spatial_vector_component_constructor(
         input: wp.array(dtype=wptype),
@@ -140,19 +140,19 @@ def test_spatial_vector_indexing(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)]
-    spatial_vector = wp.types.vector(length=6, dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    spatial_vector = wp._src.types.vector(length=6, dtype=wptype)
 
     def check_spatial_vector_indexing(
         input: wp.array(dtype=spatial_vector),
         out: wp.array(dtype=wptype),
     ):
-        inpt = input[0]
+        elem = input[0]
 
         # multiply outputs by 2 so we've got something to backpropagate:
         idx = 0
         for i in range(6):
-            out[idx] = wptype(2) * inpt[i]
+            out[idx] = wptype(2) * elem[i]
             idx = idx + 1
 
     kernel = getkernel(check_spatial_vector_indexing, suffix=dtype.__name__)
@@ -192,8 +192,8 @@ def test_spatial_vector_scalar_multiplication(test, device, dtype, register_kern
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    spatial_vector = wp.types.vector(length=6, dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    spatial_vector = wp._src.types.vector(length=6, dtype=wptype)
 
     def check_spatial_vector_scalar_mul(
         s: wp.array(dtype=wptype),
@@ -262,8 +262,8 @@ def test_spatial_vector_add_sub(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)]
-    spatial_vector = wp.types.vector(length=6, dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    spatial_vector = wp._src.types.vector(length=6, dtype=wptype)
 
     def check_spatial_vector_add_sub(
         q: wp.array(dtype=spatial_vector),
@@ -338,8 +338,8 @@ def test_spatial_dot(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)]
-    spatial_vector = wp.types.vector(length=6, dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    spatial_vector = wp._src.types.vector(length=6, dtype=wptype)
 
     def check_spatial_dot(
         s: wp.array(dtype=spatial_vector),
@@ -390,8 +390,8 @@ def test_spatial_cross(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)]
-    spatial_vector = wp.types.vector(length=6, dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    spatial_vector = wp._src.types.vector(length=6, dtype=wptype)
 
     def check_spatial_cross(
         s: wp.array(dtype=spatial_vector),
@@ -544,8 +544,8 @@ def test_spatial_top_bottom(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)]
-    spatial_vector = wp.types.vector(length=6, dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    spatial_vector = wp._src.types.vector(length=6, dtype=wptype)
 
     def check_spatial_top_bottom(
         s: wp.array(dtype=spatial_vector),
@@ -613,25 +613,25 @@ def test_transform_constructors(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)
-    transform = wp.types.transformation(dtype=wptype)
-    quat = wp.types.quaternion(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
+    transform = wp._src.types.transformation(dtype=wptype)
+    quat = wp._src.types.quaternion(dtype=wptype)
 
     def check_transform_constructor(
         input: wp.array(dtype=wptype),
         out: wp.array(dtype=wptype),
     ):
-        result = transform(vec3(input[0], input[1], input[2]), quat(input[3], input[4], input[5], input[6]))
+        a = transform(vec3(input[0], input[1], input[2]), quat(input[3], input[4], input[5], input[6]))
+        b = transform(input[0], input[1], input[2], input[3], input[4], input[5], input[6])
 
-        # multiply the output by 2 so we've got something to backpropagate:
-        out[0] = wptype(2) * result[0]
-        out[1] = wptype(2) * result[1]
-        out[2] = wptype(2) * result[2]
-        out[3] = wptype(2) * result[3]
-        out[4] = wptype(2) * result[4]
-        out[5] = wptype(2) * result[5]
-        out[6] = wptype(2) * result[6]
+        out[0] = a[0] + b[0]
+        out[1] = a[1] + b[1]
+        out[2] = a[2] + b[2]
+        out[3] = a[3] + b[3]
+        out[4] = a[4] + b[4]
+        out[5] = a[5] + b[5]
+        out[6] = a[6] + b[6]
 
     kernel = getkernel(check_transform_constructor, suffix=dtype.__name__)
     output_select_kernel = get_select_kernel(wptype)
@@ -672,19 +672,19 @@ def test_transform_indexing(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)]
-    transform = wp.types.transformation(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    transform = wp._src.types.transformation(dtype=wptype)
 
     def check_transform_indexing(
         input: wp.array(dtype=transform),
         out: wp.array(dtype=wptype),
     ):
-        inpt = input[0]
+        elem = input[0]
 
         # multiply outputs by 2 so we've got something to backpropagate:
         idx = 0
         for i in range(7):
-            out[idx] = wptype(2) * inpt[i]
+            out[idx] = wptype(2) * elem[i]
             idx = idx + 1
 
     kernel = getkernel(check_transform_indexing, suffix=dtype.__name__)
@@ -721,8 +721,8 @@ def test_transform_scalar_multiplication(test, device, dtype, register_kernels=F
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    transform = wp.types.transformation(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    transform = wp._src.types.transformation(dtype=wptype)
 
     def check_transform_scalar_mul(
         s: wp.array(dtype=wptype),
@@ -789,8 +789,8 @@ def test_transform_add_sub(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)]
-    transform = wp.types.transformation(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    transform = wp._src.types.transformation(dtype=wptype)
 
     def check_transform_add_sub(
         q: wp.array(dtype=transform),
@@ -866,8 +866,8 @@ def test_transform_get_trans_rot(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)]
-    transform = wp.types.transformation(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    transform = wp._src.types.transformation(dtype=wptype)
 
     def check_transform_get_trans_rot(
         s: wp.array(dtype=transform),
@@ -936,8 +936,8 @@ def test_transform_multiply(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)]
-    transform = wp.types.transformation(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    transform = wp._src.types.transformation(dtype=wptype)
 
     def check_transform_multiply(
         a: wp.array(dtype=transform),
@@ -1045,8 +1045,8 @@ def test_transform_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)]
-    transform = wp.types.transformation(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    transform = wp._src.types.transformation(dtype=wptype)
 
     def check_transform_inverse(
         a: wp.array(dtype=transform),
@@ -1138,9 +1138,9 @@ def test_transform_point_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)]
-    transform = wp.types.transformation(dtype=wptype)
-    vec3 = wp.types.vector(length=3, dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    transform = wp._src.types.transformation(dtype=wptype)
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
 
     def check_transform_point_vector(
         t: wp.array(dtype=transform),
@@ -1246,8 +1246,8 @@ def test_spatial_matrix_constructors(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)]
-    spatial_matrix = wp.types.matrix(shape=(6, 6), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    spatial_matrix = wp._src.types.matrix(shape=(6, 6), dtype=wptype)
 
     def check_spatial_matrix_constructor(
         input: wp.array(dtype=wptype),
@@ -1342,20 +1342,20 @@ def test_spatial_matrix_indexing(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)]
-    spatial_matrix = wp.types.matrix(shape=(6, 6), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    spatial_matrix = wp._src.types.matrix(shape=(6, 6), dtype=wptype)
 
     def check_spatial_matrix_indexing(
         input: wp.array(dtype=spatial_matrix),
         out: wp.array(dtype=wptype),
     ):
-        inpt = input[0]
+        elem = input[0]
 
         # multiply outputs by 2 so we've got something to backpropagate:
         idx = 0
         for i in range(6):
             for j in range(6):
-                out[idx] = wptype(2) * inpt[i, j]
+                out[idx] = wptype(2) * elem[i, j]
                 idx = idx + 1
 
     kernel = getkernel(check_spatial_matrix_indexing, suffix=dtype.__name__)
@@ -1397,8 +1397,8 @@ def test_spatial_matrix_scalar_multiplication(test, device, dtype, register_kern
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    spatial_matrix = wp.types.matrix(shape=(6, 6), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    spatial_matrix = wp._src.types.matrix(shape=(6, 6), dtype=wptype)
 
     def check_spatial_matrix_scalar_mul(
         s: wp.array(dtype=wptype),
@@ -1473,8 +1473,8 @@ def test_spatial_matrix_add_sub(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)]
-    spatial_matrix = wp.types.matrix(shape=(6, 6), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    spatial_matrix = wp._src.types.matrix(shape=(6, 6), dtype=wptype)
 
     def check_spatial_matrix_add_sub(
         q: wp.array(dtype=spatial_matrix),
@@ -1561,9 +1561,9 @@ def test_spatial_matvec_multiplication(test, device, dtype, register_kernels=Fal
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    spatial_matrix = wp.types.matrix(shape=(6, 6), dtype=wptype)
-    spatial_vector = wp.types.vector(length=6, dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    spatial_matrix = wp._src.types.matrix(shape=(6, 6), dtype=wptype)
+    spatial_vector = wp._src.types.vector(length=6, dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -1622,8 +1622,8 @@ def test_spatial_matmat_multiplication(test, device, dtype, register_kernels=Fal
         np.float64: 5.0e-7,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    spatial_matrix = wp.types.matrix(shape=(6, 6), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    spatial_matrix = wp._src.types.matrix(shape=(6, 6), dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -1689,8 +1689,8 @@ def test_spatial_mat_transpose(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)]
-    spatial_matrix = wp.types.matrix(shape=(6, 6), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    spatial_matrix = wp._src.types.matrix(shape=(6, 6), dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -1746,8 +1746,8 @@ def test_spatial_outer_product(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)]
-    spatial_vector = wp.types.vector(length=6, dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    spatial_vector = wp._src.types.vector(length=6, dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -1826,8 +1826,8 @@ def test_spatial_adjoint(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)]
-    mat3 = wp.types.matrix(shape=(3, 3), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    mat3 = wp._src.types.matrix(shape=(3, 3), dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -1901,7 +1901,7 @@ def test_spatial_adjoint(test, device, dtype, register_kernels=False):
 
 
 def test_transform_identity(test, device, dtype, register_kernels=False):
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
 
     def transform_identity_test(output: wp.array(dtype=wptype)):
         t = wp.transform_identity(dtype=wptype)
@@ -1936,7 +1936,7 @@ def test_transform_identity(test, device, dtype, register_kernels=False):
 def test_transform_anon_type_instance(test, device, dtype, register_kernels=False):
     rng = np.random.default_rng(123)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
 
     def transform_create_test(input: wp.array(dtype=wptype), output: wp.array(dtype=wptype)):
         t = wp.transformation(
@@ -1970,10 +1970,10 @@ def test_transform_anon_type_instance(test, device, dtype, register_kernels=Fals
 
 
 def test_transform_from_matrix(test, device, dtype, register_kernels=False):
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat44 = wp.types.matrix((4, 4), wptype)
-    vec3 = wp.types.vector(3, wptype)
-    quat = wp.types.quaternion(wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    mat44 = wp._src.types.matrix((4, 4), wptype)
+    vec3 = wp._src.types.vector(3, wptype)
+    quat = wp._src.types.quaternion(wptype)
 
     def transform_from_matrix_kernel():
         # fmt: off
@@ -1999,10 +1999,10 @@ def test_transform_from_matrix(test, device, dtype, register_kernels=False):
 
 
 def test_transform_to_matrix(test, device, dtype, register_kernels=False):
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat44 = wp.types.matrix((4, 4), wptype)
-    vec3 = wp.types.vector(3, wptype)
-    quat = wp.types.quaternion(wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    mat44 = wp._src.types.matrix((4, 4), wptype)
+    vec3 = wp._src.types.vector(3, wptype)
+    quat = wp._src.types.quaternion(wptype)
 
     def transform_to_matrix_kernel():
         p = vec3(wptype(1.0), wptype(2.0), wptype(3.0))
@@ -2031,10 +2031,10 @@ def test_transform_to_matrix(test, device, dtype, register_kernels=False):
 
 
 def test_transform_compose(test, device, dtype, register_kernels=False):
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat44 = wp.types.matrix((4, 4), wptype)
-    vec3 = wp.types.vector(3, wptype)
-    quat = wp.types.quaternion(wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    mat44 = wp._src.types.matrix((4, 4), wptype)
+    vec3 = wp._src.types.vector(3, wptype)
+    quat = wp._src.types.quaternion(wptype)
 
     def transform_compose_kernel():
         p = vec3(wptype(1.0), wptype(2.0), wptype(3.0))
@@ -2063,10 +2063,10 @@ def test_transform_compose(test, device, dtype, register_kernels=False):
 
 
 def test_transform_decompose(test, device, dtype, register_kernels=False):
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat44 = wp.types.matrix((4, 4), wptype)
-    vec3 = wp.types.vector(3, wptype)
-    quat = wp.types.quaternion(wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    mat44 = wp._src.types.matrix((4, 4), wptype)
+    vec3 = wp._src.types.vector(3, wptype)
+    quat = wp._src.types.quaternion(wptype)
 
     def transform_decompose_kernel():
         # fmt: off
@@ -2519,7 +2519,6 @@ def test_transform_slicing_assign(test, device):
     vec2 = wp.vec(2, float)
     vec3 = wp.vec(3, float)
     vec4 = wp.vec(4, float)
-    vec5 = wp.vec(5, float)
     vec6 = wp.vec(6, float)
     vec7 = wp.vec(7, float)
 
@@ -2612,6 +2611,24 @@ def test_transform_slicing_assign_backward(test, device):
     assert_np_equal(x.grad.numpy(), np.array(((1.0, 1.0),), dtype=float))
 
 
+def test_transform_default_q_arg(test, device):
+    vec7 = wp.vec(7, float)
+
+    @wp.func
+    def fn():
+        t = wp.transform(p=wp.vec3(1.0, 2.0, 3.0))
+
+        wp.expect_eq(t[:] == vec7(1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0), True)
+
+    @wp.kernel(module="unique")
+    def kernel():
+        fn()
+
+    wp.launch(kernel, 1, device=device)
+    wp.synchronize()
+    fn()
+
+
 devices = get_test_devices()
 
 
@@ -2834,6 +2851,7 @@ add_function_test(TestSpatial, "test_transform_slicing_assign", test_transform_s
 add_function_test(
     TestSpatial, "test_transform_slicing_assign_backward", test_transform_slicing_assign_backward, devices=devices
 )
+add_function_test(TestSpatial, "test_transform_default_q_arg", test_transform_default_q_arg, devices=devices)
 
 
 if __name__ == "__main__":

warp/tests/test_special_values.py

@@ -152,7 +152,7 @@ def test_nan_scalar(test, device, dtype, register_kernels=False):
 
 
 def test_is_special_vec(test, device, dtype, register_kernels=False):
-    vector_type = wp.types.vector(5, dtype)
+    vector_type = wp._src.types.vector(5, dtype)
 
     def check_special_vec(bool_outputs: wp.array(dtype=wp.bool)):
         zeros_vector = vector_type()
@@ -197,7 +197,7 @@ def test_is_special_vec(test, device, dtype, register_kernels=False):
 
 
 def test_is_special_mat(test, device, dtype, register_kernels=False):
-    mat_type = wp.types.matrix((5, 5), dtype)
+    mat_type = wp._src.types.matrix((5, 5), dtype)
 
     def check_special_mat(bool_outputs: wp.array(dtype=wp.bool)):
         zeros_mat = mat_type()
@@ -242,7 +242,7 @@ def test_is_special_mat(test, device, dtype, register_kernels=False):
 
 
 def test_is_special_quat(test, device, dtype, register_kernels=False):
-    quat_type = wp.types.quaternion(dtype)
+    quat_type = wp._src.types.quaternion(dtype)
 
     def check_special_quat(bool_outputs: wp.array(dtype=wp.bool)):
         zeros_quat = quat_type()
@@ -285,8 +285,8 @@ def test_is_special_quat(test, device, dtype, register_kernels=False):
 
 
 def test_is_special_int(test, device, dtype, register_kernels=False):
-    vector_type = wp.types.vector(5, dtype)
-    matrix_type = wp.types.matrix((5, 5), dtype)
+    vector_type = wp._src.types.vector(5, dtype)
+    matrix_type = wp._src.types.matrix((5, 5), dtype)
 
     def check_is_special_int(bool_outputs: wp.array(dtype=wp.bool)):
         bool_outputs[0] = wp.isfinite(dtype(0))
@@ -350,7 +350,7 @@ for dtype in [wp.float16, wp.float32, wp.float64]:
     )
 
 # Ensure functions like wp.isfinite work on integer types
-for dtype in wp.types.int_types:
+for dtype in wp._src.types.int_types:
     add_function_test_register_kernel(
         TestSpecialValues, f"test_is_special_int_{dtype.__name__}", test_is_special_int, devices=devices, dtype=dtype
     )

warp/tests/test_static.py

@@ -155,7 +155,7 @@ def test_invalid_static_expression(test, device):
         wp.static(1.0 / 0.0)
 
     with test.assertRaisesRegex(
-        warp.codegen.WarpCodegenError, r"Error evaluating static expression\: float division by zero"
+        warp._src.codegen.WarpCodegenError, r"Error evaluating static expression\: (?:float )?division by zero"
     ):
         wp.launch(invalid_kernel, 1, device=device)
 
@@ -164,7 +164,7 @@ def test_invalid_static_expression(test, device):
         wp.static(i * 2)
 
     with test.assertRaisesRegex(
-        wp.codegen.WarpCodegenError,
+        wp._src.codegen.WarpCodegenError,
         r"Error evaluating static expression\: name 'i' is not defined\. Make sure all variables used in the static expression are constant\.",
     ):
         wp.launch(invalid_kernel, 1, device=device, inputs=[3])
@@ -176,7 +176,7 @@ def test_static_expression_return_types(test, device):
         wp.static(wp.zeros(3, device=device))
 
     with test.assertRaisesRegex(
-        warp.codegen.WarpCodegenError,
+        warp._src.codegen.WarpCodegenError,
         r"Static expression returns an unsupported value\: a Warp array cannot be created inside Warp kernels",
     ):
         wp.launch(invalid_kernel, 1, device=device)
@@ -211,7 +211,7 @@ def test_static_expression_return_types(test, device):
         wp.static(create_struct())
 
     with test.assertRaisesRegex(
-        warp.codegen.WarpCodegenError,
+        warp._src.codegen.WarpCodegenError,
         r"Static expression returns an unsupported value: the returned Warp struct contains a data type that cannot be constructed inside Warp kernels\: a Warp array cannot be created inside Warp kernels at .*?Foo\.bar\.baz",
     ):
         wp.launch(invalid_kernel, 1, device=device)
@@ -224,7 +224,7 @@ def test_static_expression_return_types(test, device):
         wp.static(function_with_no_return_value())
 
     with test.assertRaisesRegex(
-        warp.codegen.WarpCodegenError,
+        warp._src.codegen.WarpCodegenError,
         r"Static expression returns an unsupported value\: None is returned",
     ):
         wp.launch(invalid_kernel, 1, device=device)
@@ -237,7 +237,7 @@ def test_static_expression_return_types(test, device):
         wp.static(MyClass())
 
     with test.assertRaisesRegex(
-        warp.codegen.WarpCodegenError,
+        warp._src.codegen.WarpCodegenError,
         r"Static expression returns an unsupported value\: value of type .*?MyClass",
     ):
         wp.launch(invalid_kernel, 1, device=device)
@@ -489,7 +489,7 @@ def test_static_constant_hash(test, _):
             test.assertEqual(hash1, hash3)
 
     # Warp types (scalars, vectors, matrices)
-    for warp_type in [*wp.types.scalar_types, *wp.types.vector_types]:
+    for warp_type in [*wp._src.types.scalar_types, *wp._src.types.vector_types]:
         type_name = warp_type.__name__
         with test.subTest(msg=f"wp.{type_name}"):
             value1 = ", ".join([str(17)] * warp_type._length_)

warp/tests/test_struct.py

@@ -243,12 +243,18 @@ class MatStruct:
 @wp.kernel
 def kernel_nested_mat(out: wp.array(dtype=wp.mat44)):
     s = MatStruct()
-    m = wp.mat44()
 
+    s.m[0, 0] = 2.0
     s.m[1, 2] = 3.0
+    s.m[2][1] = 5.0
 
     out[0] = s.m
 
+    out[0][2, 2] = 6.0
+    out[0][1][1] = 7.0
+
+    out[0][3, 3] = out[0][0][0]
+
 
 def test_nested_mat(test, device):
     m = wp.array([wp.mat44()], dtype=wp.mat44, device=device)
@@ -256,7 +262,12 @@ def test_nested_mat(test, device):
     wp.synchronize()
 
     out = m.numpy()
+    assert_np_equal(out[0][0, 0], 2.0)
     assert_np_equal(out[0][1, 2], 3.0)
+    assert_np_equal(out[0][2][1], 5.0)
+    assert_np_equal(out[0][2, 2], 6.0)
+    assert_np_equal(out[0][1][1], 7.0)
+    assert_np_equal(out[0][3, 3], 2.0)
 
 
 def test_assign_view(test, device):
@@ -267,7 +278,7 @@ def test_assign_view(test, device):
     m = wp.array([[wp.mat44()]], dtype=wp.mat44, device=device)
 
     with test.assertRaisesRegex(
-        wp.codegen.WarpCodegenError,
+        wp._src.codegen.WarpCodegenError,
         r"Incorrect number of indices specified for array indexing",
     ):
         wp.launch(kernel_assign_view, dim=[1, 1], outputs=[m], device=device)

warp/tests/test_triangle_closest_point.py

@@ -15,10 +15,57 @@
 
 import unittest
 
-from warp.sim.collide import triangle_closest_point_barycentric
 from warp.tests.unittest_utils import *
 
 
+@wp.func
+def triangle_closest_point_barycentric(a: wp.vec3, b: wp.vec3, c: wp.vec3, p: wp.vec3):
+    ab = b - a
+    ac = c - a
+    ap = p - a
+
+    d1 = wp.dot(ab, ap)
+    d2 = wp.dot(ac, ap)
+
+    if d1 <= 0.0 and d2 <= 0.0:
+        return wp.vec3(1.0, 0.0, 0.0)
+
+    bp = p - b
+    d3 = wp.dot(ab, bp)
+    d4 = wp.dot(ac, bp)
+
+    if d3 >= 0.0 and d4 <= d3:
+        return wp.vec3(0.0, 1.0, 0.0)
+
+    vc = d1 * d4 - d3 * d2
+    v = d1 / (d1 - d3)
+    if vc <= 0.0 and d1 >= 0.0 and d3 <= 0.0:
+        return wp.vec3(1.0 - v, v, 0.0)
+
+    cp = p - c
+    d5 = wp.dot(ab, cp)
+    d6 = wp.dot(ac, cp)
+
+    if d6 >= 0.0 and d5 <= d6:
+        return wp.vec3(0.0, 0.0, 1.0)
+
+    vb = d5 * d2 - d1 * d6
+    w = d2 / (d2 - d6)
+    if vb <= 0.0 and d2 >= 0.0 and d6 <= 0.0:
+        return wp.vec3(1.0 - w, 0.0, w)
+
+    va = d3 * d6 - d5 * d4
+    w = (d4 - d3) / ((d4 - d3) + (d5 - d6))
+    if va <= 0.0 and (d4 - d3) >= 0.0 and (d5 - d6) >= 0.0:
+        return wp.vec3(0.0, 1.0 - w, w)
+
+    denom = 1.0 / (va + vb + vc)
+    v = vb * denom
+    w = vc * denom
+
+    return wp.vec3(1.0 - v - w, v, w)
+
+
 # a-b is the edge where the closest point is located at
 @wp.func
 def check_edge_feasible_region(p: wp.vec3, a: wp.vec3, b: wp.vec3, c: wp.vec3, eps: float):