warp-lang 1.4.2__py3-none-manylinux2014_x86_64.whl → 1.5.1__py3-none-manylinux2014_x86_64.whl

warp/tests/aux_test_instancing_gc.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Helper module used in test_codegen_instancing.py"""
+
+import warp as wp
+
+
+def create_kernel_closure(value: int):
+    @wp.kernel
+    def k(a: wp.array(dtype=int)):
+        a[0] = value
+
+    return k

warp/tests/test_array.py

@@ -6,6 +6,7 @@
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
 import unittest
+from typing import Any
 
 import numpy as np
 
@@ -2429,6 +2430,16 @@ def inplace_add_rhs(x: wp.array(dtype=float), y: wp.array(dtype=float), z: wp.ar
     wp.atomic_add(z, 0, a)
 
 
+vec9 = wp.vec(length=9, dtype=float)
+
+
+@wp.kernel
+def inplace_add_custom_vec(x: wp.array(dtype=vec9), y: wp.array(dtype=vec9)):
+    i = wp.tid()
+    x[i] += y[i]
+    x[i] += y[i]
+
+
 def test_array_inplace_diff_ops(test, device):
     N = 3
     x1 = wp.ones(N, dtype=float, requires_grad=True, device=device)
@@ -2537,6 +2548,18 @@ def test_array_inplace_diff_ops(test, device):
 
     assert_np_equal(x.grad.numpy(), np.ones(1, dtype=float))
     assert_np_equal(y.grad.numpy(), np.ones(1, dtype=float))
+    tape.reset()
+
+    x = wp.zeros(1, dtype=vec9, requires_grad=True, device=device)
+    y = wp.ones(1, dtype=vec9, requires_grad=True, device=device)
+
+    with tape:
+        wp.launch(inplace_add_custom_vec, 1, inputs=[x, y], device=device)
+
+    tape.backward(grads={x: wp.ones_like(x)})
+
+    assert_np_equal(x.numpy(), np.full((1, 9), 2.0, dtype=float))
+    assert_np_equal(y.grad.numpy(), np.full((1, 9), 2.0, dtype=float))
 
 
 @wp.kernel
@@ -2551,6 +2574,15 @@ def inplace_div_1d(x: wp.array(dtype=float), y: wp.array(dtype=float)):
     x[i] /= y[i]
 
 
+@wp.kernel
+def inplace_add_non_atomic_types(x: wp.array(dtype=Any), y: wp.array(dtype=Any)):
+    i = wp.tid()
+    x[i] += y[i]
+
+
+uint16vec3 = wp.vec(length=3, dtype=wp.uint16)
+
+
 def test_array_inplace_non_diff_ops(test, device):
     N = 3
     x1 = wp.full(N, value=10.0, dtype=float, device=device)
@@ -2564,6 +2596,13 @@ def test_array_inplace_non_diff_ops(test, device):
     wp.launch(inplace_div_1d, N, inputs=[x1, y1], device=device)
     assert_np_equal(x1.numpy(), np.full(N, fill_value=2.0, dtype=float))
 
+    for dtype in wp.types.non_atomic_types + (wp.vec2b, wp.vec2ub, wp.vec2s, wp.vec2us, uint16vec3):
+        x = wp.full(N, value=0, dtype=dtype, device=device)
+        y = wp.full(N, value=1, dtype=dtype, device=device)
+
+        wp.launch(inplace_add_non_atomic_types, N, inputs=[x, y], device=device)
+        assert_np_equal(x.numpy(), y.numpy())
+
 
 @wp.kernel
 def inc_scalar(a: wp.array(dtype=float)):

warp/tests/test_codegen.py

@@ -7,11 +7,27 @@
 
 import sys
 import unittest
+from typing import Tuple
 
 import warp as wp
 from warp.tests.unittest_utils import *
 
 
+@wp.kernel
+def test_expect():
+    a = 1.0
+    a += 2.0
+
+    wp.expect_eq(123, 123)
+    wp.expect_neq(123, 234)
+
+    wp.expect_eq(wp.vec2(1.0, 2.0), wp.vec2(1.0, 2.0))
+    wp.expect_neq(wp.vec2(1.0, 2.0), wp.vec2(2.0, 3.0))
+
+    wp.expect_eq(wp.mat22(1.0, 2.0, 3.0, 4.0), wp.mat22(1.0, 2.0, 3.0, 4.0))
+    wp.expect_neq(wp.mat22(1.0, 2.0, 3.0, 4.0), wp.mat22(2.0, 3.0, 4.0, 5.0))
+
+
 @wp.kernel
 def test_rename():
     a = 0
@@ -574,6 +590,64 @@ def test_error_mutating_constant_in_dynamic_loop(test, device):
     assert_np_equal(mats.numpy(), np.zeros((1, 3, 3)))
 
 
+def test_error_return_annotation_mismatch(test, device):
+    @wp.func
+    def foo_1(x: wp.int32) -> wp.int16:
+        return wp.int8(x)
+
+    def kernel_1_fn():
+        x = foo_1(123)
+
+    @wp.func
+    def foo_2(x: int) -> int:
+        return (x + x, x * x)
+
+    def kernel_2_fn():
+        x = foo_2(123)
+
+    @wp.func
+    def foo_3(x: int) -> Tuple[int, int]:
+        return (x, 1.23)
+
+    def kernel_3_fn():
+        x, y = foo_3(123)
+
+    @wp.func
+    def foo_4(x: int) -> Tuple[int, int, int]:
+        return (x + x, x * x)
+
+    def kernel_4_fn():
+        x, y, z = foo_4(123)
+
+    kernel = wp.Kernel(func=kernel_1_fn)
+    with test.assertRaisesRegex(
+        wp.codegen.WarpCodegenError,
+        r"The function `foo_1` has its return type annotated as `int16` but the code returns a value of type `int8`.",
+    ):
+        wp.launch(kernel, dim=1, device=device)
+
+    kernel = wp.Kernel(func=kernel_2_fn)
+    with test.assertRaisesRegex(
+        wp.codegen.WarpCodegenError,
+        r"The function `foo_2` has its return type annotated as `int` but the code returns 2 values.",
+    ):
+        wp.launch(kernel, dim=1, device=device)
+
+    kernel = wp.Kernel(func=kernel_3_fn)
+    with test.assertRaisesRegex(
+        wp.codegen.WarpCodegenError,
+        r"The function `foo_3` has its return type annotated as `Tuple\[int, int\]` but the code returns a tuple with types `\(int32, float32\)`.",
+    ):
+        wp.launch(kernel, dim=1, device=device)
+
+    kernel = wp.Kernel(func=kernel_4_fn)
+    with test.assertRaisesRegex(
+        wp.codegen.WarpCodegenError,
+        r"The function `foo_4` has its return type annotated as a tuple of 3 elements but the code returns 2 values.",
+    ):
+        wp.launch(kernel, dim=1, device=device)
+
+
 @wp.kernel
 def test_call_syntax():
     expected_pow = 16.0
@@ -622,6 +696,7 @@ class TestCodeGen(unittest.TestCase):
 
 devices = get_test_devices()
 
+add_kernel_test(TestCodeGen, name="test_expect", kernel=test_expect, dim=1, devices=devices)
 add_kernel_test(TestCodeGen, name="test_inplace", kernel=test_inplace, dim=1, devices=devices)
 add_kernel_test(TestCodeGen, name="test_rename", kernel=test_rename, dim=1, devices=devices)
 add_kernel_test(TestCodeGen, name="test_constant", kernel=test_constant, inputs=[1.0], dim=1, devices=devices)
@@ -758,7 +833,12 @@ add_function_test(
     name="test_error_mutating_constant_in_dynamic_loop",
     devices=devices,
 )
-
+add_function_test(
+    TestCodeGen,
+    func=test_error_return_annotation_mismatch,
+    name="test_error_return_annotation_mismatch",
+    devices=devices,
+)
 add_kernel_test(TestCodeGen, name="test_call_syntax", kernel=test_call_syntax, dim=1, devices=devices)
 add_kernel_test(TestCodeGen, name="test_shadow_builtin", kernel=test_shadow_builtin, dim=1, devices=devices)
 add_kernel_test(TestCodeGen, name="test_while_condition_eval", kernel=test_while_condition_eval, dim=1, devices=devices)

warp/tests/test_codegen_instancing.py

@@ -1287,6 +1287,35 @@ def test_module_mark_modified(test, device):
 # =======================================================================
 
 
+def test_garbage_collection(test, device):
+    """Test that dynamically generated kernels without user references are not retained in the module."""
+
+    # use a helper module with a known kernel count
+    import warp.tests.aux_test_instancing_gc as gc_test_module
+
+    with wp.ScopedDevice(device):
+        a = wp.zeros(1, dtype=int)
+
+        for i in range(10):
+            # create a unique kernel on each iteration
+            k = gc_test_module.create_kernel_closure(i)
+
+            # import gc
+            # gc.collect()
+
+            # since we don't keep references to the previous kernels,
+            # they should be garbage-collected and not appear in the module
+            k.module.load(device=device)
+            test.assertEqual(len(k.module.live_kernels), 1)
+
+            # test the kernel
+            wp.launch(k, dim=1, inputs=[a])
+            test.assertEqual(a.numpy()[0], i)
+
+
+# =======================================================================
+
+
 class TestCodeGenInstancing(unittest.TestCase):
     pass
 
@@ -1450,6 +1479,7 @@ add_function_test(TestCodeGenInstancing, func=test_create_kernel_loop, name="tes
 add_function_test(
     TestCodeGenInstancing, func=test_module_mark_modified, name="test_module_mark_modified", devices=devices
 )
+add_function_test(TestCodeGenInstancing, func=test_garbage_collection, name="test_garbage_collection", devices=devices)
 
 
 if __name__ == "__main__":

warp/tests/test_collision.py

@@ -7,8 +7,11 @@
 
 import unittest
 
+import warp as wp
 import warp.examples
+import warp.sim
 from warp.sim.collide import *
+from warp.sim.integrator_euler import eval_triangles_contact
 from warp.tests.unittest_utils import *
 
 
@@ -474,12 +477,119 @@ def test_edge_edge_collision(test, device):
 devices = get_test_devices()
 
 
+def test_particle_collision(test, device):
+    with wp.ScopedDevice(device):
+        contact_radius = 1.23
+        builder1 = wp.sim.ModelBuilder()
+        builder1.add_cloth_grid(
+            pos=wp.vec3(0.0, 0.0, 0.0),
+            rot=wp.quat_identity(),
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            dim_x=100,
+            dim_y=100,
+            cell_x=0.1,
+            cell_y=0.1,
+            mass=0.1,
+            particle_radius=contact_radius,
+        )
+
+    cloth_grid = builder1.finalize()
+    cloth_grid_particle_radius = cloth_grid.particle_radius.numpy()
+    assert_np_equal(cloth_grid_particle_radius, np.full(cloth_grid_particle_radius.shape, contact_radius), tol=1e-5)
+
+    vertices = [
+        [2.0, 0.0, 0.0],
+        [2.0, 2.0, 0.0],
+        [0.0, 0.0, 0.0],
+        [1.0, 0.0, 1.0],
+        [1.0, 1.0, 1.0],
+        [0.0, 0.0, 1.0],
+    ]
+    vertices = [wp.vec3(v) for v in vertices]
+    faces = [0, 1, 2, 3, 4, 5]
+
+    builder2 = wp.sim.ModelBuilder()
+    builder2.add_cloth_mesh(
+        pos=wp.vec3(0.0, 0.0, 0.0),
+        rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
+        scale=1.0,
+        vertices=vertices,
+        indices=faces,
+        tri_ke=1e4,
+        tri_ka=1e4,
+        tri_kd=1e-5,
+        edge_ke=10,
+        edge_kd=0.0,
+        vel=wp.vec3(0.0, 0.0, 0.0),
+        density=0.1,
+        particle_radius=contact_radius,
+    )
+    cloth_mesh = builder2.finalize()
+    cloth_mesh_particle_radius = cloth_mesh.particle_radius.numpy()
+    assert_np_equal(cloth_mesh_particle_radius, np.full(cloth_mesh_particle_radius.shape, contact_radius), tol=1e-5)
+
+    state = cloth_mesh.state()
+    particle_f = wp.zeros_like(state.particle_q)
+    wp.launch(
+        kernel=eval_triangles_contact,
+        dim=cloth_mesh.tri_count * cloth_mesh.particle_count,
+        inputs=[
+            cloth_mesh.particle_count,
+            state.particle_q,
+            state.particle_qd,
+            cloth_mesh.tri_indices,
+            cloth_mesh.tri_materials,
+            cloth_mesh.particle_radius,
+        ],
+        outputs=[particle_f],
+    )
+    test.assertTrue((np.linalg.norm(particle_f.numpy(), axis=1) != 0).all())
+
+    builder3 = wp.sim.ModelBuilder()
+    builder3.add_cloth_mesh(
+        pos=wp.vec3(0.0, 0.0, 0.0),
+        rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
+        scale=1.0,
+        vertices=vertices,
+        indices=faces,
+        tri_ke=1e4,
+        tri_ka=1e4,
+        tri_kd=1e-5,
+        edge_ke=10,
+        edge_kd=0.0,
+        vel=wp.vec3(0.0, 0.0, 0.0),
+        density=0.1,
+        particle_radius=0.5,
+    )
+    cloth_mesh_2 = builder3.finalize()
+    cloth_mesh_2_particle_radius = cloth_mesh_2.particle_radius.numpy()
+    assert_np_equal(cloth_mesh_2_particle_radius, np.full(cloth_mesh_2_particle_radius.shape, 0.5), tol=1e-5)
+
+    state_2 = cloth_mesh_2.state()
+    particle_f_2 = wp.zeros_like(cloth_mesh_2.particle_q)
+    wp.launch(
+        kernel=eval_triangles_contact,
+        dim=cloth_mesh_2.tri_count * cloth_mesh_2.particle_count,
+        inputs=[
+            cloth_mesh_2.particle_count,
+            state_2.particle_q,
+            state_2.particle_qd,
+            cloth_mesh_2.tri_indices,
+            cloth_mesh_2.tri_materials,
+            cloth_mesh_2.particle_radius,
+        ],
+        outputs=[particle_f_2],
+    )
+    test.assertTrue((np.linalg.norm(particle_f_2.numpy(), axis=1) == 0).all())
+
+
 class TestCollision(unittest.TestCase):
     pass
 
 
 add_function_test(TestCollision, "test_vertex_triangle_collision", test_vertex_triangle_collision, devices=devices)
 add_function_test(TestCollision, "test_edge_edge_collision", test_vertex_triangle_collision, devices=devices)
+add_function_test(TestCollision, "test_particle_collision", test_particle_collision, devices=devices)
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/test_coloring.py

@@ -0,0 +1,251 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+
+import numpy as np
+
+import warp as wp
+import warp.examples
+import warp.sim
+from warp.sim.graph_coloring import (
+    ColoringAlgorithm,
+    construct_trimesh_graph_edges,
+    convert_to_color_groups,
+    validate_graph_coloring,
+)
+from warp.tests.unittest_utils import *
+
+
+def color_lattice_grid(num_x, num_y):
+    colors = []
+    for _ in range(4):
+        colors.append([])
+
+    for xi in range(num_x + 1):
+        for yi in range(num_y + 1):
+            node_dx = yi * (num_x + 1) + xi
+
+            a = 1 if xi % 2 else 0
+            b = 1 if yi % 2 else 0
+
+            c = b * 2 + a
+
+            colors[c].append(node_dx)
+
+    color_groups = [np.array(group) for group in colors]
+
+    return color_groups
+
+
+@unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
+def test_coloring_trimesh(test, device):
+    from pxr import Usd, UsdGeom
+
+    with wp.ScopedDevice(device):
+        usd_stage = Usd.Stage.Open(os.path.join(wp.examples.get_asset_directory(), "bunny.usd"))
+        usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/bunny"))
+
+        vertices = np.array(usd_geom.GetPointsAttr().Get())
+        faces = np.array(usd_geom.GetFaceVertexIndicesAttr().Get())
+
+        builder = wp.sim.ModelBuilder()
+
+        builder.add_cloth_mesh(
+            pos=wp.vec3(0.0, 0.0, 0.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
+            scale=1.0,
+            vertices=[wp.vec3(p) for p in vertices],
+            indices=faces.flatten(),
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            density=0.02,
+        )
+
+        model = builder.finalize()
+
+        particle_colors = wp.empty(shape=(model.particle_count), dtype=int, device="cpu")
+
+        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(
+            model.particle_count,
+            edge_indices_cpu.__ctype__(),
+            ColoringAlgorithm.GREEDY.value,
+            particle_colors.__ctype__(),
+        )
+        wp.launch(
+            kernel=validate_graph_coloring,
+            inputs=[edge_indices_cpu, particle_colors],
+            dim=edge_indices_cpu.shape[0],
+            device="cpu",
+        )
+
+        num_colors_mcs = wp.context.runtime.core.graph_coloring(
+            model.particle_count,
+            edge_indices_cpu.__ctype__(),
+            ColoringAlgorithm.MCS.value,
+            particle_colors.__ctype__(),
+        )
+        wp.launch(
+            kernel=validate_graph_coloring,
+            inputs=[edge_indices_cpu, particle_colors],
+            dim=edge_indices_cpu.shape[0],
+            device="cpu",
+        )
+
+        # 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(
+            model.particle_count,
+            edge_indices_cpu_with_bending.__ctype__(),
+            ColoringAlgorithm.GREEDY.value,
+            particle_colors.__ctype__(),
+        )
+        wp.context.runtime.core.balance_coloring(
+            model.particle_count,
+            edge_indices_cpu_with_bending.__ctype__(),
+            num_colors_greedy,
+            1.1,
+            particle_colors.__ctype__(),
+        )
+        wp.launch(
+            kernel=validate_graph_coloring,
+            inputs=[edge_indices_cpu_with_bending, particle_colors],
+            dim=edge_indices_cpu_with_bending.shape[0],
+            device="cpu",
+        )
+
+        num_colors_mcs = wp.context.runtime.core.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(
+            model.particle_count,
+            edge_indices_cpu_with_bending.__ctype__(),
+            num_colors_mcs,
+            1.1,
+            particle_colors.__ctype__(),
+        )
+        wp.launch(
+            kernel=validate_graph_coloring,
+            inputs=[edge_indices_cpu_with_bending, particle_colors],
+            dim=edge_indices_cpu_with_bending.shape[0],
+            device="cpu",
+        )
+
+        color_categories_balanced = convert_to_color_groups(num_colors_mcs, particle_colors)
+
+        color_sizes = np.array([c.shape[0] for c in color_categories_balanced], dtype=np.float32)
+        test.assertTrue(np.max(color_sizes) / np.min(color_sizes) <= max_min_ratio)
+
+
+@unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
+def test_combine_coloring(test, device):
+    from pxr import Usd, UsdGeom
+
+    with wp.ScopedDevice(device):
+        builder1 = wp.sim.ModelBuilder()
+        usd_stage = Usd.Stage.Open(os.path.join(wp.examples.get_asset_directory(), "bunny.usd"))
+        usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/bunny"))
+
+        vertices = np.array(usd_geom.GetPointsAttr().Get())
+        faces = np.array(usd_geom.GetFaceVertexIndicesAttr().Get())
+
+        builder1.add_cloth_mesh(
+            pos=wp.vec3(0.0, 0.0, 0.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
+            scale=1.0,
+            vertices=[wp.vec3(p) for p in vertices],
+            indices=faces.flatten(),
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            density=0.02,
+        )
+
+        builder1.add_cloth_grid(
+            pos=wp.vec3(0.0, 4.0, 0.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            dim_x=50,
+            dim_y=100,
+            cell_x=0.1,
+            cell_y=0.1,
+            mass=0.1,
+            fix_left=True,
+        )
+        builder1.color()
+
+        builder2 = wp.sim.ModelBuilder()
+        builder2.add_cloth_grid(
+            pos=wp.vec3(0.0, 4.0, 0.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            dim_x=50,
+            dim_y=100,
+            cell_x=0.1,
+            cell_y=0.1,
+            mass=0.1,
+            # to include bending in coloring
+            edge_ke=100000,
+            fix_left=True,
+        )
+        builder2.color()
+
+        builder3 = wp.sim.ModelBuilder()
+        builder3.add_cloth_grid(
+            pos=wp.vec3(0.0, 4.0, 0.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            dim_x=50,
+            dim_y=100,
+            cell_x=0.1,
+            cell_y=0.1,
+            mass=0.1,
+            fix_left=True,
+        )
+
+        builder3.set_coloring(
+            color_lattice_grid(50, 100),
+        )
+
+        builder1.add_builder(builder2)
+        builder1.add_builder(builder3)
+
+        model = builder2.finalize()
+
+        particle_number_colored = np.full((model.particle_count), -1, dtype=int)
+        particle_colors = np.full((model.particle_count), -1, dtype=int)
+        for color, color_group in enumerate(model.particle_coloring):
+            particle_number_colored[color_group.numpy()] += 1
+            particle_colors[color_group.numpy()] = color
+
+        # all particles has been colored exactly once
+        assert_np_equal(particle_number_colored, 0)
+
+        edge_indices_cpu = wp.array(model.edge_indices.numpy()[:, 2:], dtype=int, device="cpu")
+        wp.launch(
+            kernel=validate_graph_coloring,
+            inputs=[edge_indices_cpu, wp.array(particle_colors, dtype=int, device="cpu")],
+            dim=edge_indices_cpu.shape[0],
+            device="cpu",
+        )
+
+
+devices = get_test_devices()
+
+
+class TestColoring(unittest.TestCase):
+    pass
+
+
+add_function_test(TestColoring, "test_coloring_trimesh", test_coloring_trimesh, devices=devices)
+add_function_test(TestColoring, "test_combine_coloring", test_combine_coloring, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_context.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+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)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)