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_mat_scalar_ops.py

@@ -75,13 +75,13 @@ def get_select_kernel(dtype):
 def test_arrays(test, device, dtype):
     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)]
 
-    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)
-    mat32 = wp.types.matrix(shape=(3, 2), dtype=wptype)
+    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)
+    mat32 = wp._src.types.matrix(shape=(3, 2), dtype=wptype)
 
     v2_np = randvals(rng, [10, 2, 2], dtype)
     v3_np = randvals(rng, [10, 3, 3], dtype)
@@ -101,9 +101,9 @@ def test_arrays(test, device, dtype):
     assert_np_equal(v5.numpy(), v5_np, tol=1.0e-6)
     assert_np_equal(v32.numpy(), v32_np, tol=1.0e-6)
 
-    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)
+    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)
 
     v2 = wp.array(v2_np, dtype=mat22, requires_grad=True, device=device)
     v3 = wp.array(v3_np, dtype=mat33, requires_grad=True, device=device)
@@ -118,8 +118,8 @@ def test_components(test, device, dtype):
     # test accessing matrix components from Python - this is especially important
     # for float16, which requires special handling internally
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat23 = wp.types.matrix(shape=(2, 3), dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    mat23 = wp._src.types.matrix(shape=(2, 3), dtype=wptype)
 
     m = mat23(1, 2, 3, 4, 5, 6)
 
@@ -167,12 +167,12 @@ def test_components(test, device, dtype):
 
 
 def test_constants(test, device, dtype, register_kernels=False):
-    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)
-    mat32 = wp.types.matrix(shape=(3, 2), 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)
+    mat32 = wp._src.types.matrix(shape=(3, 2), dtype=wptype)
 
     cm22 = wp.constant(mat22(22))
     cm33 = wp.constant(mat33(33))
@@ -202,15 +202,15 @@ def test_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)]
-    vec2 = wp.types.vector(length=2, dtype=wptype)
-    vec3 = wp.types.vector(length=3, dtype=wptype)
-    vec4 = wp.types.vector(length=4, dtype=wptype)
-    vec5 = wp.types.vector(length=5, dtype=wptype)
-    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)]
+    vec2 = wp._src.types.vector(length=2, dtype=wptype)
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
+    vec4 = wp._src.types.vector(length=4, dtype=wptype)
+    vec5 = wp._src.types.vector(length=5, dtype=wptype)
+    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)
 
@@ -446,7 +446,7 @@ def test_anon_type_instance(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)]
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
 
     def check_scalar_init(
         input: wp.array(dtype=wptype),
@@ -629,7 +629,7 @@ def test_anon_type_instance(test, device, dtype, register_kernels=False):
 
 
 def test_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 check_identity_mat(
         output: wp.array(dtype=wptype),
@@ -679,11 +679,11 @@ def test_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)]
-    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)
 
@@ -755,11 +755,11 @@ def test_indexing(test, device, dtype, register_kernels=False):
 
 
 def test_equality(test, device, dtype, register_kernels=False):
-    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)
 
     def check_mat_equality():
         wp.expect_eq(
@@ -1029,11 +1029,11 @@ def test_scalar_multiplication(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)
 
@@ -1216,17 +1216,17 @@ def test_matvec_multiplication(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)
-    mat32 = wp.types.matrix(shape=(3, 2), 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)
+    mat32 = wp._src.types.matrix(shape=(3, 2), dtype=wptype)
+    mat44 = wp._src.types.matrix(shape=(4, 4), dtype=wptype)
+    mat55 = wp._src.types.matrix(shape=(5, 5), dtype=wptype)
 
-    vec2 = wp.types.vector(length=2, dtype=wptype)
-    vec3 = wp.types.vector(length=3, dtype=wptype)
-    vec4 = wp.types.vector(length=4, dtype=wptype)
-    vec5 = wp.types.vector(length=5, dtype=wptype)
+    vec2 = wp._src.types.vector(length=2, dtype=wptype)
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
+    vec4 = wp._src.types.vector(length=4, dtype=wptype)
+    vec5 = wp._src.types.vector(length=5, dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -1363,17 +1363,17 @@ def test_vecmat_multiplication(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)
-    mat23 = wp.types.matrix(shape=(2, 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)
+    mat23 = wp._src.types.matrix(shape=(2, 3), dtype=wptype)
+    mat44 = wp._src.types.matrix(shape=(4, 4), dtype=wptype)
+    mat55 = wp._src.types.matrix(shape=(5, 5), dtype=wptype)
 
-    vec2 = wp.types.vector(length=2, dtype=wptype)
-    vec3 = wp.types.vector(length=3, dtype=wptype)
-    vec4 = wp.types.vector(length=4, dtype=wptype)
-    vec5 = wp.types.vector(length=5, dtype=wptype)
+    vec2 = wp._src.types.vector(length=2, dtype=wptype)
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
+    vec4 = wp._src.types.vector(length=4, dtype=wptype)
+    vec5 = wp._src.types.vector(length=5, dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -1512,12 +1512,12 @@ def test_matmat_multiplication(test, device, dtype, register_kernels=False):
         np.float64: 5.0e-7,
     }.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)
-    mat32 = wp.types.matrix(shape=(3, 2), 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)
+    mat32 = wp._src.types.matrix(shape=(3, 2), 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)
 
@@ -1691,11 +1691,11 @@ def test_cw_multiplication(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)
 
@@ -1823,11 +1823,11 @@ def test_cw_division(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)
 
@@ -1980,11 +1980,11 @@ def test_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)]
-    vec2 = wp.types.vector(length=2, dtype=wptype)
-    vec3 = wp.types.vector(length=3, dtype=wptype)
-    vec4 = wp.types.vector(length=4, dtype=wptype)
-    vec5 = wp.types.vector(length=5, dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    vec2 = wp._src.types.vector(length=2, dtype=wptype)
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
+    vec4 = wp._src.types.vector(length=4, dtype=wptype)
+    vec5 = wp._src.types.vector(length=5, dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -2120,12 +2120,12 @@ def test_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)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat32 = wp.types.matrix(shape=(3, 2), 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)
+    mat32 = wp._src.types.matrix(shape=(3, 2), 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)
 
@@ -2219,11 +2219,11 @@ def test_scalar_division(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)
 
@@ -2321,11 +2321,11 @@ def test_addition(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)
 
@@ -2451,11 +2451,11 @@ def test_ddot(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)
 
     def check_mat_dot(
         s2: wp.array(dtype=mat22),
@@ -2574,11 +2574,11 @@ def test_trace(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)
 
     def check_mat_trace(
         v2: wp.array(dtype=mat22),
@@ -2666,8 +2666,8 @@ def test_diag(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)]
-    vec5 = wp.types.vector(length=5, dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    vec5 = wp._src.types.vector(length=5, dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -2716,19 +2716,19 @@ def test_diag(test, device, dtype, register_kernels=False):
 
 
 def test_equivalent_types(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)]
 
     # matrix types
-    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)
+    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)
 
     # matrix types equivalent to the above
-    mat22_equiv = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33_equiv = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44_equiv = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55_equiv = wp.types.matrix(shape=(5, 5), dtype=wptype)
+    mat22_equiv = wp._src.types.matrix(shape=(2, 2), dtype=wptype)
+    mat33_equiv = wp._src.types.matrix(shape=(3, 3), dtype=wptype)
+    mat44_equiv = wp._src.types.matrix(shape=(4, 4), dtype=wptype)
+    mat55_equiv = wp._src.types.matrix(shape=(5, 5), dtype=wptype)
 
     # declare kernel with original types
     def check_equivalence(

warp/tests/test_math.py

@@ -58,6 +58,39 @@ def test_scalar_math(test, device):
         assert_np_equal(tape.gradients[x].numpy(), np.array([adj_float_results_expected[i]]), tol=1e-6)
 
 
+@wp.kernel
+def erf_kernel(x: wp.array(dtype=Any), out: wp.array(dtype=Any)):
+    i = wp.tid()
+
+    if i == 0:
+        out[i] = wp.erf(x[i])
+    elif i == 1:
+        out[i] = wp.erfc(x[i])
+    elif i == 2:
+        out[i] = wp.erfinv(x[i])
+    elif i == 3:
+        out[i] = wp.erfcinv(x[i])
+
+
+def test_erf_math(test, device):
+    for type, tol in ((wp.float16, 1e-3), (wp.float32, 1e-6), (wp.float64, 1e-6)):
+        x = wp.full(4, value=0.123, dtype=type, requires_grad=True, device=device)
+        out = wp.zeros(4, dtype=type, requires_grad=True, device=device)
+
+        with wp.Tape() as tape:
+            wp.launch(erf_kernel, dim=4, inputs=[x], outputs=[out], device=device)
+
+        out.grad = wp.ones_like(out)
+
+        tape.backward()
+
+        out_true = np.array([0.13809388, 0.86190612, 0.10944129, 1.09057285])
+        adj_x_true = np.array([1.11143641, -1.11143641, 0.89690544, -2.91120449])
+
+        assert_np_equal(out.numpy(), out_true, tol=tol)
+        assert_np_equal(adj_x_true, x.grad.numpy(), tol=tol)
+
+
 @wp.kernel
 def test_vec_norm_kernel(vs: wp.array(dtype=Any), out: wp.array(dtype=float, ndim=2)):
     tid = wp.tid()
@@ -170,6 +203,7 @@ class TestMath(unittest.TestCase):
 
 
 add_function_test(TestMath, "test_scalar_math", test_scalar_math, devices=devices)
+add_function_test(TestMath, "test_erf_math", test_erf_math, devices=devices)
 add_function_test(TestMath, "test_vec_norm", test_vec_norm, devices=devices)
 
 

warp/tests/test_module_aot.py

@@ -226,14 +226,14 @@ class TestModuleAOT(unittest.TestCase):
         if wp.get_cuda_device_count() == 0:
             self.skipTest("No CUDA devices found")
 
-        if len(wp.context.runtime.nvrtc_supported_archs) < 2:
+        if len(wp._src.context.runtime.nvrtc_supported_archs) < 2:
             self.skipTest("NVRTC must support at least two architectures to run this test")
 
         try:
             shutil.rmtree(TEST_CACHE_DIR, ignore_errors=True)
             TEST_CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
-            archs = list(wp.context.runtime.nvrtc_supported_archs)[:2]
+            archs = list(wp._src.context.runtime.nvrtc_supported_archs)[:2]
 
             wp.compile_aot_module(warp.tests.aux_test_module_aot, arch=archs, module_dir=TEST_CACHE_DIR, use_ptx=True)
 
@@ -254,14 +254,14 @@ class TestModuleAOT(unittest.TestCase):
         if wp.get_cuda_device_count() == 0:
             self.skipTest("No CUDA devices found")
 
-        if len(wp.context.runtime.nvrtc_supported_archs) < 2:
+        if len(wp._src.context.runtime.nvrtc_supported_archs) < 2:
             self.skipTest("NVRTC must support at least two architectures to run this test")
 
         try:
             shutil.rmtree(TEST_CACHE_DIR, ignore_errors=True)
             TEST_CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
-            archs = list(wp.context.runtime.nvrtc_supported_archs)[:2]
+            archs = list(wp._src.context.runtime.nvrtc_supported_archs)[:2]
 
             wp.compile_aot_module(warp.tests.aux_test_module_aot, arch=archs, module_dir=TEST_CACHE_DIR, use_ptx=False)
 

warp/tests/test_modules_lite.py

@@ -16,21 +16,30 @@
 import unittest
 
 import warp as wp
+import warp._src.utils
 from warp.tests.unittest_utils import *
 
 devices = get_test_devices()
 
 
+# This kernel is needed to ensure this test module is registered as a Warp module.
+# wp.load_module() requires the module to contain at least one Warp kernel, function, or struct.
+@wp.kernel
+def print_values():
+    i = wp.tid()
+    wp.print(i)
+
+
 class TestModuleLite(unittest.TestCase):
     def test_module_lite_load(self):
         # Load current module
         wp.load_module()
 
         # Load named module
-        wp.load_module(wp.config)
+        wp.load_module(warp._src.utils)
 
         # Load named module (string)
-        wp.load_module(wp.config, recursive=True)
+        wp.load_module("warp._src.utils", recursive=True)
 
     def test_module_lite_options(self):
         wp.set_module_options({"max_unroll": 8})
@@ -38,6 +47,23 @@ class TestModuleLite(unittest.TestCase):
         self.assertIsInstance(module_options, dict)
         self.assertEqual(module_options["max_unroll"], 8)
 
+    def test_module_lite_load_nonexistent(self):
+        # Test that loading a non-existent module raises RuntimeError
+        with self.assertRaises(RuntimeError) as context:
+            wp.load_module("nonexistent_module_that_does_not_exist")
+
+        self.assertIn("does not contain any Warp kernels, functions, or structs", str(context.exception))
+        self.assertIn("nonexistent_module_that_does_not_exist", str(context.exception))
+
+    def test_module_lite_load_no_warp_content(self):
+        # Test that loading a module without Warp content raises RuntimeError
+        # Use a standard library module that definitely has no Warp kernels
+        with self.assertRaises(RuntimeError) as context:
+            wp.load_module(unittest)
+
+        self.assertIn("does not contain any Warp kernels, functions, or structs", str(context.exception))
+        self.assertIn("unittest", str(context.exception))
+
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/test_print.py

@@ -234,14 +234,14 @@ class SimpleStruct:
     y: float
 
 
-generic_print_types = [*wp.types.scalar_types]
-for scalar_type in wp.types.scalar_types:
-    generic_print_types.append(wp.types.vector(2, scalar_type))
-    generic_print_types.append(wp.types.vector(3, scalar_type))
-    generic_print_types.append(wp.types.vector(4, scalar_type))
-    generic_print_types.append(wp.types.matrix((2, 2), scalar_type))
-    generic_print_types.append(wp.types.matrix((3, 3), scalar_type))
-    generic_print_types.append(wp.types.matrix((4, 4), scalar_type))
+generic_print_types = [*wp._src.types.scalar_types]
+for scalar_type in wp._src.types.scalar_types:
+    generic_print_types.append(wp._src.types.vector(2, scalar_type))
+    generic_print_types.append(wp._src.types.vector(3, scalar_type))
+    generic_print_types.append(wp._src.types.vector(4, scalar_type))
+    generic_print_types.append(wp._src.types.matrix((2, 2), scalar_type))
+    generic_print_types.append(wp._src.types.matrix((3, 3), scalar_type))
+    generic_print_types.append(wp._src.types.matrix((4, 4), scalar_type))
 generic_print_types.append(wp.bool)
 generic_print_types.append(SimpleStruct)
 generic_print_types.append(wp.array(dtype=float))
@@ -251,7 +251,7 @@ for T in generic_print_types:
 
 
 def test_print_adjoint(test, device):
-    for scalar_type in wp.types.scalar_types:
+    for scalar_type in wp._src.types.scalar_types:
         # scalar
         capture = StdOutCapture()
         capture.begin()
@@ -272,7 +272,7 @@ def test_print_adjoint(test, device):
 
         for dim in (2, 3, 4):
             # vector
-            vec_type = wp.types.vector(dim, scalar_type)
+            vec_type = wp._src.types.vector(dim, scalar_type)
             vec_data = np.arange(vec_type._length_, dtype=wp.dtype_to_numpy(scalar_type))
             v = vec_type(vec_data)
             adj_v = vec_type(vec_data[::-1])
@@ -290,7 +290,7 @@ def test_print_adjoint(test, device):
                 test.assertRegex(s, rf"{expected_forward}{os.linesep}adj: {expected_adjoint}{os.linesep}")
 
             # matrix
-            mat_type = wp.types.matrix((dim, dim), scalar_type)
+            mat_type = wp._src.types.matrix((dim, dim), scalar_type)
             mat_data = np.arange(mat_type._length_, dtype=wp.dtype_to_numpy(scalar_type))
             m = mat_type(mat_data)
             adj_m = mat_type(mat_data[::-1])