warp-lang 1.8.1__py3-none-win_amd64.whl → 1.9.0__py3-none-win_amd64.whl

warp/tests/aux_test_module_aot.py

@@ -0,0 +1,7 @@
+import warp as wp
+
+
+@wp.kernel
+def add_kernel(a: wp.array(dtype=wp.int32), b: wp.array(dtype=wp.int32), res: wp.array(dtype=wp.int32)):
+    i = wp.tid()
+    res[i] = a[i] + b[i]

warp/tests/cuda/test_async.py

@@ -426,13 +426,13 @@ def copy_template(test, src_ctor, dst_ctor, src_device, dst_device, n, params: C
         if expected_error_type is not None:
             # disable error output from Warp if we expect an exception
             try:
-                saved_error_output_enabled = wp.context.runtime.core.is_error_output_enabled()
-                wp.context.runtime.core.set_error_output_enabled(False)
+                saved_error_output_enabled = wp.context.runtime.core.wp_is_error_output_enabled()
+                wp.context.runtime.core.wp_set_error_output_enabled(False)
                 with test.assertRaisesRegex(expected_error_type, expected_error_regex):
                     with Capturable(use_graph=params.use_graph, stream=stream):
                         wp.copy(dst, src, stream=stream_arg)
             finally:
-                wp.context.runtime.core.set_error_output_enabled(saved_error_output_enabled)
+                wp.context.runtime.core.wp_set_error_output_enabled(saved_error_output_enabled)
                 wp.synchronize()
 
                 # print(f"SUCCESSFUL ERROR PREDICTION: {expected_error_regex}")

warp/tests/cuda/test_conditional_captures.py

@@ -979,6 +979,94 @@ def test_graph_debug_dot_print(test, device):
     os.remove(dot_file)
 
 
+# ================================================================================================================
+# test exceptions
+# ================================================================================================================
+
+
+# body graphs with allocations are not supported
+def body_with_alloc():
+    wp.zeros(10)
+
+
+@unittest.skipUnless(check_conditional_graph_support(), "Conditional graph nodes not supported")
+def test_error_alloc_if(test, device):
+    with wp.ScopedDevice(device):
+        cond = wp.ones(1, dtype=wp.int32)
+        with test.assertRaisesRegex(
+            RuntimeError, r"Conditional body graph contains an unsupported operation \(memory allocation\)"
+        ):
+            with wp.ScopedCapture():
+                wp.capture_if(condition=cond, on_true=body_with_alloc)
+
+
+@unittest.skipUnless(check_conditional_graph_support(), "Conditional graph nodes not supported")
+def test_error_alloc_else(test, device):
+    with wp.ScopedDevice(device):
+        cond = wp.ones(1, dtype=wp.int32)
+        with test.assertRaisesRegex(
+            RuntimeError, r"Conditional body graph contains an unsupported operation \(memory allocation\)"
+        ):
+            with wp.ScopedCapture():
+                wp.capture_if(condition=cond, on_false=body_with_alloc)
+
+
+@unittest.skipUnless(check_conditional_graph_support(), "Conditional graph nodes not supported")
+def test_error_alloc_while(test, device):
+    with wp.ScopedDevice(device):
+        cond = wp.ones(1, dtype=wp.int32)
+        with test.assertRaisesRegex(
+            RuntimeError, r"Conditional body graph contains an unsupported operation \(memory allocation\)"
+        ):
+            with wp.ScopedCapture():
+                wp.capture_while(condition=cond, while_body=body_with_alloc)
+
+
+@unittest.skipUnless(check_conditional_graph_support(), "Conditional graph nodes not supported")
+def test_error_alloc_if_subgraph(test, device):
+    with wp.ScopedDevice(device):
+        # capture body subgraph
+        with wp.ScopedCapture() as body_capture:
+            body_with_alloc()
+
+        cond = wp.ones(1, dtype=wp.int32)
+        with test.assertRaisesRegex(
+            RuntimeError, r"Child graph contains an unsupported operation \(memory allocation\)"
+        ):
+            with wp.ScopedCapture():
+                wp.capture_if(condition=cond, on_true=body_capture.graph)
+
+
+@unittest.skipUnless(check_conditional_graph_support(), "Conditional graph nodes not supported")
+def test_error_alloc_else_subgraph(test, device):
+    with wp.ScopedDevice(device):
+        # capture body subgraph
+        with wp.ScopedCapture() as body_capture:
+            body_with_alloc()
+
+        cond = wp.ones(1, dtype=wp.int32)
+        with test.assertRaisesRegex(
+            RuntimeError, r"Child graph contains an unsupported operation \(memory allocation\)"
+        ):
+            with wp.ScopedCapture():
+                wp.capture_if(condition=cond, on_false=body_capture.graph)
+
+
+@unittest.skipUnless(check_conditional_graph_support(), "Conditional graph nodes not supported")
+def test_error_alloc_while_subgraph(test, device):
+    with wp.ScopedDevice(device):
+        # capture body subgraph
+        with wp.ScopedCapture() as body_capture:
+            body_with_alloc()
+
+        cond = wp.ones(1, dtype=wp.int32)
+        with test.assertRaisesRegex(
+            RuntimeError, r"Child graph contains an unsupported operation \(memory allocation\)"
+        ):
+            with wp.ScopedCapture():
+                wp.capture_while(condition=cond, while_body=body_capture.graph)
+
+
 devices = get_test_devices()
 cuda_devices = get_cuda_test_devices()
 
@@ -1040,6 +1128,19 @@ add_function_test(
     TestConditionalCaptures, "test_graph_debug_dot_print", test_graph_debug_dot_print, devices=cuda_devices
 )
 
+add_function_test(TestConditionalCaptures, "test_error_alloc_if", test_error_alloc_if, devices=cuda_devices)
+add_function_test(TestConditionalCaptures, "test_error_alloc_else", test_error_alloc_else, devices=cuda_devices)
+add_function_test(TestConditionalCaptures, "test_error_alloc_while", test_error_alloc_while, devices=cuda_devices)
+add_function_test(
+    TestConditionalCaptures, "test_error_alloc_if_subgraph", test_error_alloc_if_subgraph, devices=cuda_devices
+)
+add_function_test(
+    TestConditionalCaptures, "test_error_alloc_else_subgraph", test_error_alloc_else_subgraph, devices=cuda_devices
+)
+add_function_test(
+    TestConditionalCaptures, "test_error_alloc_while_subgraph", test_error_alloc_while_subgraph, devices=cuda_devices
+)
+
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/geometry/test_marching_cubes.py

@@ -22,7 +22,9 @@ from warp.tests.unittest_utils import *
 
 
 @wp.kernel
-def make_field(field: wp.array3d(dtype=float), center: wp.vec3, radius: float):
+def make_field_sphere_sdf(field: wp.array3d(dtype=float), center: wp.vec3, radius: float):
+    """Make a sphere SDF for nodes on the integer domain with node coordinates 0,1,2,3,..."""
+
     i, j, k = wp.tid()
 
     p = wp.vec3(float(i), float(j), float(k))
@@ -32,31 +34,242 @@ def make_field(field: wp.array3d(dtype=float), center: wp.vec3, radius: float):
     field[i, j, k] = d
 
 
-def test_marching_cubes(test, device):
-    dim = 64
-    max_verts = 10**6
-    max_tris = 10**6
+@wp.kernel
+def make_field_sphere_sdf_unit_domain(
+    field: wp.array3d(dtype=float), center: wp.vec3, radius: wp.array(dtype=wp.float32)
+):
+    """Makes a sphere SDF for nodes on the unit domain [-1, 1]^3."""
+    i, j, k = wp.tid()
 
-    field = wp.zeros(shape=(dim, dim, dim), dtype=float, device=device)
+    nx, ny, nz = field.shape[0], field.shape[1], field.shape[2]
+
+    p = wp.vec3(
+        2.0 * wp.float32(i) / (wp.float32(nx) - 1.0) - 1.0,
+        2.0 * wp.float32(j) / (wp.float32(ny) - 1.0) - 1.0,
+        2.0 * wp.float32(k) / (wp.float32(nz) - 1.0) - 1.0,
+    )
+
+    d = wp.length(p - center) - radius[0]
+
+    field[i, j, k] = d
+
+
+@wp.kernel
+def compute_surface_area(
+    verts: wp.array(dtype=wp.vec3), faces: wp.array(dtype=wp.int32), out_area: wp.array(dtype=wp.float32)
+):
+    tid = wp.tid()
+    vi = faces[3 * tid + 0]
+    vj = faces[3 * tid + 1]
+    vk = faces[3 * tid + 2]
+
+    p0 = verts[vi]
+    p1 = verts[vj]
+    p2 = verts[vk]
+
+    # Heron's formula for triangle area
+    a = wp.length(p1 - p0)
+    b = wp.length(p2 - p0)
+    c = wp.length(p2 - p1)
+    s = (a + b + c) / 2.0
+    area = wp.sqrt(s * (s - a) * (s - b) * (s - c))
 
-    iso = wp.MarchingCubes(nx=dim, ny=dim, nz=dim, max_verts=max_verts, max_tris=max_tris, device=device)
+    wp.atomic_add(out_area, 0, area)
 
-    radius = dim / 4.0
 
-    wp.launch(make_field, dim=field.shape, inputs=[field, wp.vec3(dim / 2, dim / 2, dim / 2), radius], device=device)
+def validate_marching_cubes_output(test, verts_np, faces_np, check_nonempty=True):
+    # check that the face array seems valid
+    if check_nonempty:
+        test.assertGreater(faces_np.shape[0], 0)  # at least one face
+    test.assertEqual(faces_np.shape[1], 3)  # all faces triangular
+    test.assertTrue((faces_np >= 0).all())  # all face inds nonnegative
+    test.assertTrue((faces_np < verts_np.shape[0]).all())  # all face inds are in-bounds on the vertex array
+    test.assertTrue((faces_np[:, 0] != faces_np[:, 1]).all())  # all faces have unique vertices
+    test.assertTrue((faces_np[:, 0] != faces_np[:, 2]).all())  # all faces have unique vertices
+    test.assertTrue((faces_np[:, 1] != faces_np[:, 2]).all())  # all faces have unique vertices
+    test.assertTrue(
+        (np.unique(faces_np.flatten()) == np.arange(verts_np.shape[0])).all()
+    )  # all vertices are used in at least one face
+
+    # check that the vertex array seems valid
+    if check_nonempty:
+        test.assertGreater(verts_np.shape[0], 0)  # at least one vertex
+    test.assertEqual(verts_np.shape[1], 3)  # all vertices are 3D
+    test.assertTrue(np.isfinite(verts_np).all())  # all vertices are finite
+
+
+def test_marching_cubes(test, device):
+    """Basic test of typical usage."""
+    node_dim = 64
+    cell_dim = node_dim - 1
+    field = wp.zeros(shape=(node_dim, node_dim, node_dim), dtype=float, device=device)
+    bounds_low = (0.0, 0.0, 0.0)
+    bounds_high = (float(cell_dim), float(cell_dim), float(cell_dim))
+
+    iso = wp.MarchingCubes(
+        nx=node_dim,
+        ny=node_dim,
+        nz=node_dim,
+        device=device,
+        domain_bounds_lower_corner=bounds_low,
+        domain_bounds_upper_corner=bounds_high,
+    )
+
+    radius = node_dim / 4.0
+
+    wp.launch(
+        make_field_sphere_sdf,
+        dim=field.shape,
+        inputs=[field, wp.vec3(node_dim / 2, node_dim / 2, node_dim / 2), radius],
+        device=device,
+    )
 
     iso.surface(field=field, threshold=0.0)
+    verts_np = iso.verts.numpy()
+    faces_np = iso.indices.numpy().reshape(-1, 3)
+    test.assertEqual(
+        iso.indices.dtype, wp.int32
+    )  # make sure we are following Warp convention of using a flat array of indices
+    validate_marching_cubes_output(test, verts_np, faces_np)
 
     # check that all returned vertices lie on the surface of the sphere
-    length = np.linalg.norm(iso.verts.numpy() - np.array([dim / 2, dim / 2, dim / 2]), axis=1)
+    length = np.linalg.norm(verts_np - np.array([node_dim / 2, node_dim / 2, node_dim / 2]), axis=1)
     error = np.abs(length - radius)
+    test.assertTrue(np.max(error) < 1.0)
+
+    iso.resize(nx=node_dim * 2, ny=node_dim * 2, nz=node_dim * 2)  # smoke test for deprecated function
+
+
+def test_marching_cubes_functional(test, device):
+    """Ensure the single-function interface works as expected."""
+    node_dim = 64
+    cell_dim = node_dim - 1
+    field = wp.zeros(shape=(node_dim, node_dim, node_dim), dtype=float, device=device)
+    bounds_low = (0.0, 0.0, 0.0)
+    bounds_high = (float(cell_dim), float(cell_dim), float(cell_dim))
+
+    radius = node_dim / 4.0
+    wp.launch(
+        make_field_sphere_sdf,
+        dim=field.shape,
+        inputs=[field, wp.vec3(node_dim / 2, node_dim / 2, node_dim / 2), radius],
+        device=device,
+    )
+
+    # call via the functional interface
+    verts, faces = wp.MarchingCubes.extract_surface_marching_cubes(
+        field, threshold=0.0, domain_bounds_lower_corner=bounds_low, domain_bounds_upper_corner=bounds_high
+    )
+
+    verts_np = verts.numpy()
+    faces_np = faces.numpy().reshape(-1, 3)
+    validate_marching_cubes_output(test, verts_np, faces_np)
 
+    # check that all returned vertices lie on the surface of the sphere
+    length = np.linalg.norm(verts_np - np.array([node_dim / 2, node_dim / 2, node_dim / 2]), axis=1)
+    error = np.abs(length - radius)
     test.assertTrue(np.max(error) < 1.0)
 
-    iso.resize(nx=dim * 2, ny=dim * 2, nz=dim * 2, max_verts=max_verts, max_tris=max_tris)
+
+def test_marching_cubes_nonuniform(test, device):
+    """Test the logic for when the dimensions of the grid are not uniform."""
+
+    dimX = 64
+    dimY = 48
+    dimZ = 72
+    field = wp.zeros(shape=(dimX, dimY, dimZ), dtype=float, device=device)
+
+    bounds_low = wp.vec3(0.0, 0.0, 0.0)
+    bounds_high = wp.vec3(float(dimX), float(dimY), float(dimZ))
+
+    iso = wp.MarchingCubes(
+        nx=dimX,
+        ny=dimY,
+        nz=dimZ,
+        device=device,
+        domain_bounds_lower_corner=bounds_low,
+        domain_bounds_upper_corner=bounds_high,
+    )
+
+    radius = dimX / 4.0
+    wp.launch(
+        make_field_sphere_sdf,
+        dim=field.shape,
+        inputs=[field, wp.vec3(dimX / 2, dimY / 2, dimZ / 2), radius],
+        device=device,
+    )
+
+    iso.surface(field=field, threshold=0.0)
+    verts_np = iso.verts.numpy()
+    faces_np = iso.indices.numpy().reshape(-1, 3)
+    validate_marching_cubes_output(test, verts_np, faces_np)
+
+
+def test_marching_cubes_empty_output(test, device):
+    """Make sure we handle the empty-output case correctly."""
+
+    dim = 64
+    field = wp.zeros(shape=(dim, dim, dim), dtype=float, device=device)
+
+    iso = wp.MarchingCubes(nx=dim, ny=dim, nz=dim, device=device)
+
+    wp.launch(make_field_sphere_sdf, dim=field.shape, inputs=[field, wp.vec3(0.5, 0.5, 0.5), 0.25], device=device)
+
+    iso.surface(field=field, threshold=1000.0)  # set threshold to a large value so that no vertices are generated
+    verts_np = iso.verts.numpy()
+    faces_np = iso.indices.numpy().reshape(-1, 3)
+    validate_marching_cubes_output(test, verts_np, faces_np, check_nonempty=False)
+
+    test.assertEqual(faces_np.shape[0], 0)  # no faces
+    test.assertEqual(verts_np.shape[0], 0)  # no vertices
+
+
+def test_marching_cubes_differentiable(test, device):
+    """Check that marching cubes has reasonable gradients.
+
+    This test constructs an SDF of sphere, extracts a surface, computes its
+    surface area, and then differentiates the surface area with respect to
+    the sphere's radius.
+    """
+    node_dim = 64
+    cell_dim = node_dim - 1
+    bounds_low = wp.vec3(-1.0, -1.0, -1.0)
+    bounds_high = wp.vec3(1.0, 1.0, 1.0)
+
+    radius = 0.5
+    radius_wp = wp.full((1,), value=0.5, dtype=wp.float32, device=device, requires_grad=True)
+
+    with wp.Tape() as tape:
+        field = wp.zeros(shape=(node_dim, node_dim, node_dim), dtype=float, device=device, requires_grad=True)
+        wp.launch(
+            make_field_sphere_sdf_unit_domain,
+            dim=field.shape,
+            inputs=[field, wp.vec3(0.0, 0.0, 0.0), radius_wp],
+            device=device,
+        )
+
+        # call via the functional interface
+        verts, faces = wp.MarchingCubes.extract_surface_marching_cubes(
+            field, threshold=0.0, domain_bounds_lower_corner=bounds_low, domain_bounds_upper_corner=bounds_high
+        )
+
+        # compute surface area
+        area = wp.zeros(shape=(1,), dtype=float, device=device, requires_grad=True)
+        wp.launch(compute_surface_area, dim=faces.shape[0] // 3, inputs=[verts, faces, area], device=device)
+
+        # confirm surface area is correct vs. the analytical ground truth
+        area_np = area.numpy()[0]
+        test.assertTrue(np.abs(area_np - 4.0 * np.pi * radius * radius) < 1e-2)
+
+    # compute the gradient of the surface area with respect to the radius
+    tape.backward(area)
+
+    # confirm the gradient is correct vs. the analytical ground truth
+    grad_np = radius_wp.grad.numpy()[0]
+    test.assertTrue(np.abs(grad_np - 8.0 * np.pi * radius) < 1e-2)
 
 
-devices = get_selected_cuda_test_devices()
+devices = get_test_devices()
 
 
 class TestMarchingCubes(unittest.TestCase):
@@ -67,6 +280,14 @@ class TestMarchingCubes(unittest.TestCase):
 
 
 add_function_test(TestMarchingCubes, "test_marching_cubes", test_marching_cubes, devices=devices)
+add_function_test(TestMarchingCubes, "test_marching_cubes_functional", test_marching_cubes_functional, devices=devices)
+add_function_test(TestMarchingCubes, "test_marching_cubes_nonuniform", test_marching_cubes_nonuniform, devices=devices)
+add_function_test(
+    TestMarchingCubes, "test_marching_cubes_empty_output", test_marching_cubes_empty_output, devices=devices
+)
+add_function_test(
+    TestMarchingCubes, "test_marching_cubes_differentiable", test_marching_cubes_differentiable, devices=devices
+)
 
 
 if __name__ == "__main__":

warp/tests/sim/test_coloring.py

@@ -150,7 +150,7 @@ def test_coloring_trimesh(test, device):
         edge_indices_cpu = wp.array(model.edge_indices.numpy()[:, 2:], dtype=int, device="cpu")
 
         # coloring without bending
-        num_colors_greedy = wp.context.runtime.core.graph_coloring(
+        num_colors_greedy = wp.context.runtime.core.wp_graph_coloring(
             model.particle_count,
             edge_indices_cpu.__ctype__(),
             ColoringAlgorithm.GREEDY.value,
@@ -163,7 +163,7 @@ def test_coloring_trimesh(test, device):
             device="cpu",
         )
 
-        num_colors_mcs = wp.context.runtime.core.graph_coloring(
+        num_colors_mcs = wp.context.runtime.core.wp_graph_coloring(
             model.particle_count,
             edge_indices_cpu.__ctype__(),
             ColoringAlgorithm.MCS.value,
@@ -178,13 +178,13 @@ def test_coloring_trimesh(test, device):
 
         # coloring with bending
         edge_indices_cpu_with_bending = construct_trimesh_graph_edges(model.edge_indices, True)
-        num_colors_greedy = wp.context.runtime.core.graph_coloring(
+        num_colors_greedy = wp.context.runtime.core.wp_graph_coloring(
             model.particle_count,
             edge_indices_cpu_with_bending.__ctype__(),
             ColoringAlgorithm.GREEDY.value,
             particle_colors.__ctype__(),
         )
-        wp.context.runtime.core.balance_coloring(
+        wp.context.runtime.core.wp_balance_coloring(
             model.particle_count,
             edge_indices_cpu_with_bending.__ctype__(),
             num_colors_greedy,
@@ -198,13 +198,13 @@ def test_coloring_trimesh(test, device):
             device="cpu",
         )
 
-        num_colors_mcs = wp.context.runtime.core.graph_coloring(
+        num_colors_mcs = wp.context.runtime.core.wp_graph_coloring(
             model.particle_count,
             edge_indices_cpu_with_bending.__ctype__(),
             ColoringAlgorithm.MCS.value,
             particle_colors.__ctype__(),
         )
-        max_min_ratio = wp.context.runtime.core.balance_coloring(
+        max_min_ratio = wp.context.runtime.core.wp_balance_coloring(
             model.particle_count,
             edge_indices_cpu_with_bending.__ctype__(),
             num_colors_mcs,

warp/tests/test_array.py

@@ -2902,10 +2902,8 @@ def test_direct_from_numpy(test, device):
 
 
 @wp.kernel
-def kernel_array_from_ptr(
-    ptr: wp.uint64,
-):
-    arr = wp.array(ptr=ptr, shape=(2, 3), dtype=wp.float32)
+def kernel_array_from_ptr(arr_orig: wp.array2d(dtype=wp.float32)):
+    arr = wp.array(ptr=arr_orig.ptr, shape=(2, 3), dtype=wp.float32)
     arr[0, 0] = 1.0
     arr[0, 1] = 2.0
     arr[0, 2] = 3.0
@@ -2913,7 +2911,56 @@ def kernel_array_from_ptr(
 
 def test_kernel_array_from_ptr(test, device):
     arr = wp.zeros(shape=(2, 3), dtype=wp.float32, device=device)
-    wp.launch(kernel_array_from_ptr, dim=(1,), inputs=(arr.ptr,), device=device)
+    wp.launch(kernel_array_from_ptr, dim=(1,), inputs=(arr,), device=device)
+    assert_np_equal(arr.numpy(), np.array(((1.0, 2.0, 3.0), (0.0, 0.0, 0.0))))
+
+
+@wp.struct
+class MyStruct:
+    a: wp.float32
+    b: wp.float32
+    c: wp.float32
+
+
+@wp.kernel
+def kernel_array_from_ptr_struct(arr_orig: wp.array(dtype=MyStruct)):
+    arr = wp.array(ptr=arr_orig.ptr, shape=(2,), dtype=MyStruct)
+    arr[0].a = 1.0
+    arr[0].b = 2.0
+    arr[0].c = 3.0
+    arr[1].a = 4.0
+    arr[1].b = 5.0
+    arr[1].c = 6.0
+
+
+def test_kernel_array_from_ptr_struct(test, device):
+    arr = wp.zeros(shape=(2,), dtype=MyStruct, device=device)
+    wp.launch(kernel_array_from_ptr_struct, dim=(1,), inputs=(arr,), device=device)
+    arr_np = arr.numpy()
+    expected = np.zeros_like(arr_np)
+    expected[0] = (1.0, 2.0, 3.0)
+    expected[1] = (4.0, 5.0, 6.0)
+    assert_np_equal(arr_np, expected)
+
+
+@wp.kernel
+def kernel_array_from_ptr_variable_shape(
+    ptr: wp.uint64,
+    shape_x: int,
+    shape_y: int,
+):
+    arr = wp.array(ptr=ptr, shape=(shape_x, shape_y), dtype=wp.float32)
+    arr[0, 0] = 1.0
+    arr[0, 1] = 2.0
+    if shape_y > 2:
+        arr[0, 2] = 3.0
+
+
+def test_kernel_array_from_ptr_variable_shape(test, device):
+    arr = wp.zeros(shape=(2, 3), dtype=wp.float32, device=device)
+    wp.launch(kernel_array_from_ptr_variable_shape, dim=(1,), inputs=(arr.ptr, 2, 2), device=device)
+    assert_np_equal(arr.numpy(), np.array(((1.0, 2.0, 0.0), (0.0, 0.0, 0.0))))
+    wp.launch(kernel_array_from_ptr_variable_shape, dim=(1,), inputs=(arr.ptr, 2, 3), device=device)
     assert_np_equal(arr.numpy(), np.array(((1.0, 2.0, 3.0), (0.0, 0.0, 0.0))))
 
 
@@ -3185,6 +3232,10 @@ add_function_test(TestArray, "test_array_inplace_diff_ops", test_array_inplace_d
 add_function_test(TestArray, "test_array_inplace_non_diff_ops", test_array_inplace_non_diff_ops, devices=devices)
 add_function_test(TestArray, "test_direct_from_numpy", test_direct_from_numpy, devices=["cpu"])
 add_function_test(TestArray, "test_kernel_array_from_ptr", test_kernel_array_from_ptr, devices=devices)
+add_function_test(TestArray, "test_kernel_array_from_ptr_struct", test_kernel_array_from_ptr_struct, devices=devices)
+add_function_test(
+    TestArray, "test_kernel_array_from_ptr_variable_shape", test_kernel_array_from_ptr_variable_shape, devices=devices
+)
 
 add_function_test(TestArray, "test_array_from_int32_domain", test_array_from_int32_domain, devices=devices)
 add_function_test(TestArray, "test_array_from_int64_domain", test_array_from_int64_domain, devices=devices)

warp/tests/test_codegen.py

@@ -756,6 +756,7 @@ def test_multiple_return_values(test, device):
         test_multiple_return_values_quat_to_axis_angle_kernel,
         dim=1,
         inputs=(q, expected_axis, expected_angle),
+        device=device,
     )
 
     # fmt: off
@@ -791,9 +792,9 @@ def test_multiple_return_values(test, device):
 
     test.assertAlmostEqual(V[0][0], expected_V[0][0], places=5)
     test.assertAlmostEqual(V[0][1], expected_V[0][1], places=5)
-    test.assertAlmostEqual(V[0][2], expected_V[0][2], places=5)
+    test.assertAlmostEqual(V[0][2], expected_V[0][2], places=4)  # precision issue on ARM64 (GH-905)
     test.assertAlmostEqual(V[1][0], expected_V[1][0], places=5)
-    test.assertAlmostEqual(V[1][1], expected_V[1][1], places=5)
+    test.assertAlmostEqual(V[1][1], expected_V[1][1], places=4)  # precision issue on ARM64 (GH-905)
     test.assertAlmostEqual(V[1][2], expected_V[1][2], places=5)
     test.assertAlmostEqual(V[2][0], expected_V[2][0], places=5)
     test.assertAlmostEqual(V[2][1], expected_V[2][1], places=5)

warp/tests/test_context.py

@@ -17,24 +17,17 @@ import unittest
 from typing import List, Tuple
 
 import warp as wp
-from warp.tests.unittest_utils import *
-
-
-def test_context_type_str(test, device):
-    assert wp.context.type_str(List[int]) == "List[int]"
-    assert wp.context.type_str(List[float]) == "List[float]"
-
-    assert wp.context.type_str(Tuple[int]) == "Tuple[int]"
-    assert wp.context.type_str(Tuple[float]) == "Tuple[float]"
-    assert wp.context.type_str(Tuple[int, float]) == "Tuple[int, float]"
-    assert wp.context.type_str(Tuple[int, ...]) == "Tuple[int, ...]"
 
 
 class TestContext(unittest.TestCase):
-    pass
-
-
-add_function_test(TestContext, "test_context_type_str", test_context_type_str)
+    def test_context_type_str(self):
+        self.assertEqual(wp.context.type_str(List[int]), "List[int]")
+        self.assertEqual(wp.context.type_str(List[float]), "List[float]")
+
+        self.assertEqual(wp.context.type_str(Tuple[int]), "Tuple[int]")
+        self.assertEqual(wp.context.type_str(Tuple[float]), "Tuple[float]")
+        self.assertEqual(wp.context.type_str(Tuple[int, float]), "Tuple[int, float]")
+        self.assertEqual(wp.context.type_str(Tuple[int, ...]), "Tuple[int, ...]")
 
 
 if __name__ == "__main__":