warp-lang 1.5.0__py3-none-win_amd64.whl → 1.6.0__py3-none-win_amd64.whl

warp/tests/test_assert.py

@@ -0,0 +1,242 @@
+# 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
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+@wp.kernel
+def expect_ones(a: wp.array(dtype=int)):
+    i = wp.tid()
+
+    assert a[i] == 1
+
+
+@wp.kernel
+def expect_ones_with_msg(a: wp.array(dtype=int)):
+    i = wp.tid()
+
+    assert a[i] == 1, "Array element must be 1"
+
+
+@wp.kernel
+def expect_ones_compound(a: wp.array(dtype=int)):
+    i = wp.tid()
+
+    assert a[i] > 0 and a[i] < 2
+
+
+@wp.func
+def expect_ones_function(value: int):
+    assert value == 1, "Array element must be 1"
+
+
+@wp.kernel
+def expect_ones_call_function(a: wp.array(dtype=int)):
+    i = wp.tid()
+    expect_ones_function(a[i])
+
+
+class TestAssertRelease(unittest.TestCase):
+    """Assert test cases that are to be run with Warp in release mode."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls._saved_mode = wp.get_module_options()["mode"]
+        cls._saved_cache_kernels = wp.config.cache_kernels
+
+        wp.config.mode = "release"
+        wp.config.cache_kernels = False
+
+    @classmethod
+    def tearDownClass(cls):
+        wp.set_module_options({"mode": cls._saved_mode})
+        wp.config.cache_kernels = cls._saved_cache_kernels
+
+    def test_basic_assert_false_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.zeros(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+    def test_basic_assert_with_msg(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.zeros(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_with_msg, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+    def test_compound_assert_false_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.full(1, value=3, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_compound, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+    def test_basic_assert_false_condition_function(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.full(1, value=3, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_call_function, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+
+# NOTE: Failed assertions on CUDA devices leaves the CUDA context in an unrecoverable state,
+# so we currently do not test them.
+class TestAssertDebug(unittest.TestCase):
+    """Assert test cases that are to be run with Warp in debug mode."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls._saved_mode = wp.get_module_options()["mode"]
+        cls._saved_cache_kernels = wp.config.cache_kernels
+
+        wp.set_module_options({"mode": "debug"})
+        wp.config.cache_kernels = False
+
+    @classmethod
+    def tearDownClass(cls):
+        wp.set_module_options({"mode": cls._saved_mode})
+        wp.config.cache_kernels = cls._saved_cache_kernels
+
+    def test_basic_assert_false_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.zeros(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            # Older Windows C runtimes have a bug where stdout sometimes does not get properly flushed.
+            if output != "" or sys.platform != "win32":
+                self.assertRegex(output, r"Assertion failed: .*assert a\[i\] == 1")
+
+    def test_basic_assert_true_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.ones(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+    def test_basic_assert_with_msg(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.zeros(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_with_msg, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            # Older Windows C runtimes have a bug where stdout sometimes does not get properly flushed.
+            if output != "" or sys.platform != "win32":
+                self.assertRegex(output, r"Assertion failed: .*assert a\[i\] == 1.*Array element must be 1")
+
+    def test_compound_assert_true_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.ones(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_compound, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+    def test_compound_assert_false_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.full(1, value=3, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_compound, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            # Older Windows C runtimes have a bug where stdout sometimes does not get properly flushed.
+            if output != "" or sys.platform != "win32":
+                self.assertRegex(output, r"Assertion failed: .*assert a\[i\] > 0 and a\[i\] < 2")
+
+    def test_basic_assert_false_condition_function(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.full(1, value=3, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_call_function, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            # Older Windows C runtimes have a bug where stdout sometimes does not get properly flushed.
+            if output != "" or sys.platform != "win32":
+                self.assertRegex(output, r"Assertion failed: .*assert value == 1.*Array element must be 1")
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_codegen.py

@@ -396,48 +396,29 @@ def test_unresolved_symbol(test, device):
 
 
 def test_error_global_var(test, device):
-    arr = wp.array(
-        (1.0, 2.0, 3.0),
-        dtype=float,
-        device=device,
-    )
+    arr = wp.array((1.0, 2.0, 3.0), dtype=float, device=device)
 
-    def kernel_1_fn(
-        out: wp.array(dtype=float),
-    ):
+    def kernel_1_fn(out: wp.array(dtype=float)):
         out[0] = arr[0]
 
-    def kernel_2_fn(
-        out: wp.array(dtype=float),
-    ):
+    def kernel_2_fn(out: wp.array(dtype=float)):
         out[0] = arr
 
-    def kernel_3_fn(
-        out: wp.array(dtype=float),
-    ):
+    def kernel_3_fn(out: wp.array(dtype=float)):
         out[0] = wp.lower_bound(arr, 2.0)
 
     out = wp.empty_like(arr)
 
     kernel = wp.Kernel(func=kernel_1_fn)
-    with test.assertRaisesRegex(
-        TypeError,
-        r"Invalid external reference type: <class 'warp.types.array'>",
-    ):
+    with test.assertRaisesRegex(TypeError, r"Invalid external reference type: <class 'warp.types.array'>"):
         wp.launch(kernel, dim=out.shape, inputs=(), outputs=(out,), device=device)
 
     kernel = wp.Kernel(func=kernel_2_fn)
-    with test.assertRaisesRegex(
-        TypeError,
-        r"Invalid external reference type: <class 'warp.types.array'>",
-    ):
+    with test.assertRaisesRegex(TypeError, r"Invalid external reference type: <class 'warp.types.array'>"):
         wp.launch(kernel, dim=out.shape, inputs=(), outputs=(out,), device=device)
 
     kernel = wp.Kernel(func=kernel_3_fn)
-    with test.assertRaisesRegex(
-        TypeError,
-        r"Invalid external reference type: <class 'warp.types.array'>",
-    ):
+    with test.assertRaisesRegex(TypeError, r"Invalid external reference type: <class 'warp.types.array'>"):
         wp.launch(kernel, dim=out.shape, inputs=(), outputs=(out,), device=device)
 
 
@@ -469,16 +450,12 @@ def test_error_collection_construct(test, device):
         wp.launch(kernel, dim=1, device=device)
 
     kernel = wp.Kernel(func=kernel_3_fn)
-    with test.assertRaisesRegex(
-        RuntimeError,
-        r"Construct `ast.Dict` not supported in kernels.",
-    ):
+    with test.assertRaisesRegex(RuntimeError, r"Construct `ast.Dict` not supported in kernels."):
         wp.launch(kernel, dim=1, device=device)
 
     kernel = wp.Kernel(func=kernel_4_fn)
     with test.assertRaisesRegex(
-        RuntimeError,
-        r"Tuple constructs are not supported in kernels. Use vectors like `wp.vec3\(\)` instead.",
+        RuntimeError, r"Tuple constructs are not supported in kernels. Use vectors like `wp.vec3\(\)` instead."
     ):
         wp.launch(kernel, dim=1, device=device)
 
@@ -491,10 +468,7 @@ def test_error_unmatched_arguments(test, device):
         x = wp.dot(wp.vec2(1.0, 2.0), wp.vec2h(wp.float16(1.0), wp.float16(2.0)))
 
     kernel = wp.Kernel(func=kernel_1_fn)
-    with test.assertRaisesRegex(
-        RuntimeError,
-        r"Input types must be the same, got \['int32', 'float32'\]",
-    ):
+    with test.assertRaisesRegex(RuntimeError, r"Input types must be the same, got \['int32', 'float32'\]"):
         wp.launch(kernel, dim=1, device=device)
 
     kernel = wp.Kernel(func=kernel_2_fn)
@@ -704,12 +678,7 @@ add_kernel_test(
     TestCodeGen, name="test_dynamic_for_rename", kernel=test_dynamic_for_rename, inputs=[10], dim=1, devices=devices
 )
 add_kernel_test(
-    TestCodeGen,
-    name="test_dynamic_for_inplace",
-    kernel=test_dynamic_for_inplace,
-    inputs=[10],
-    dim=1,
-    devices=devices,
+    TestCodeGen, name="test_dynamic_for_inplace", kernel=test_dynamic_for_inplace, inputs=[10], dim=1, devices=devices
 )
 add_kernel_test(TestCodeGen, name="test_reassign", kernel=test_reassign, dim=1, devices=devices)
 add_kernel_test(
@@ -754,12 +723,7 @@ add_kernel_test(
 )
 
 add_kernel_test(
-    TestCodeGen,
-    name="test_range_static_sum",
-    kernel=test_range_static_sum,
-    dim=1,
-    expect=[10, 10, 10],
-    devices=devices,
+    TestCodeGen, name="test_range_static_sum", kernel=test_range_static_sum, dim=1, expect=[10, 10, 10], devices=devices
 )
 add_kernel_test(
     TestCodeGen,
@@ -789,20 +753,9 @@ add_kernel_test(
     devices=devices,
 )
 add_kernel_test(
-    TestCodeGen,
-    name="test_range_dynamic_nested",
-    kernel=test_range_dynamic_nested,
-    dim=1,
-    inputs=[4],
-    devices=devices,
-)
-add_kernel_test(
-    TestCodeGen,
-    name="test_range_expression",
-    kernel=test_range_expression,
-    dim=1,
-    devices=devices,
+    TestCodeGen, name="test_range_dynamic_nested", kernel=test_range_dynamic_nested, dim=1, inputs=[4], devices=devices
 )
+add_kernel_test(TestCodeGen, name="test_range_expression", kernel=test_range_expression, dim=1, devices=devices)
 
 add_kernel_test(TestCodeGen, name="test_while_zero", kernel=test_while, dim=1, inputs=[0], devices=devices)
 add_kernel_test(TestCodeGen, name="test_while_positive", kernel=test_while, dim=1, inputs=[16], devices=devices)

warp/tests/test_collision.py

@@ -430,8 +430,8 @@ class Example:
     def set_points_fixed(self, model, fixed_particles):
         if len(fixed_particles):
             flags = model.particle_flags.numpy()
-            for fixed_v_id in fixed_particles:
-                flags[fixed_v_id] = wp.uint32(int(flags[fixed_v_id]) & ~int(PARTICLE_FLAG_ACTIVE))
+            for fixed_vertex_id in fixed_particles:
+                flags[fixed_vertex_id] = wp.uint32(int(flags[fixed_vertex_id]) & ~int(PARTICLE_FLAG_ACTIVE))
 
             model.particle_flags = wp.array(flags, device=model.device)
 

warp/tests/test_coloring.py

@@ -11,7 +11,12 @@ 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, validate_graph_coloring
+from warp.sim.graph_coloring import (
+    ColoringAlgorithm,
+    construct_trimesh_graph_edges,
+    convert_to_color_groups,
+    validate_graph_coloring,
+)
 from warp.tests.unittest_utils import *
 
 
@@ -120,7 +125,7 @@ def test_coloring_trimesh(test, device):
             ColoringAlgorithm.MCS.value,
             particle_colors.__ctype__(),
         )
-        wp.context.runtime.core.balance_coloring(
+        max_min_ratio = wp.context.runtime.core.balance_coloring(
             model.particle_count,
             edge_indices_cpu_with_bending.__ctype__(),
             num_colors_mcs,
@@ -134,6 +139,11 @@ def test_coloring_trimesh(test, device):
             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):

warp/tests/test_examples.py

@@ -165,7 +165,9 @@ def add_example_test(
 
         # with wp.ScopedTimer(f"{name}_{sanitize_identifier(device)}"):
         # Run the script as a subprocess
-        result = subprocess.run(command, capture_output=True, text=True, env=env_vars, timeout=test_timeout)
+        result = subprocess.run(
+            command, capture_output=True, text=True, env=env_vars, timeout=test_timeout, check=False
+        )
 
         # Check the return code (0 is standard for success)
         test.assertEqual(
@@ -302,6 +304,15 @@ add_example_test(
     },
     test_options_cpu={"train_iters": 1, "num_frames": 30},
 )
+add_example_test(
+    TestOptimExamples,
+    name="optim.example_softbody_properties",
+    devices=test_devices,
+    test_options_cuda={
+        "train_iters": 1 if warp.context.runtime.core.is_debug_enabled() else 3,
+    },
+    test_options_cpu={"train_iters": 1},
+)
 
 
 class TestSimExamples(unittest.TestCase):

warp/tests/test_func.py

@@ -162,7 +162,7 @@ def user_func_with_defaults(a: int = 123, b: int = 234) -> int:
 
 
 @wp.kernel
-def test_user_func_with_defaults():
+def user_func_with_defaults_kernel():
     a = user_func_with_defaults()
     wp.expect_eq(a, 357)
 
@@ -179,6 +179,25 @@ def test_user_func_with_defaults():
     wp.expect_eq(e, 234)
 
 
+def test_user_func_with_defaults(test, device):
+    wp.launch(user_func_with_defaults_kernel, dim=1, device=device)
+
+    a = user_func_with_defaults()
+    assert a == 357
+
+    b = user_func_with_defaults(111)
+    assert b == 345
+
+    c = user_func_with_defaults(111, 222)
+    assert c == 333
+
+    d = user_func_with_defaults(a=111)
+    assert d == 345
+
+    e = user_func_with_defaults(b=111)
+    assert e == 234
+
+
 @wp.func
 def user_func_return_multiple_values(a: int, b: float) -> Tuple[int, float]:
     return a + a, b * b
@@ -406,9 +425,7 @@ add_function_test(TestFunc, func=test_func_closure_capture, name="test_func_clos
 add_function_test(TestFunc, func=test_multi_valued_func, name="test_multi_valued_func", devices=devices)
 add_kernel_test(TestFunc, kernel=test_func_defaults, name="test_func_defaults", dim=1, devices=devices)
 add_kernel_test(TestFunc, kernel=test_builtin_shadowing, name="test_builtin_shadowing", dim=1, devices=devices)
-add_kernel_test(
-    TestFunc, kernel=test_user_func_with_defaults, name="test_user_func_with_defaults", dim=1, devices=devices
-)
+add_function_test(TestFunc, func=test_user_func_with_defaults, name="test_user_func_with_defaults", devices=devices)
 add_kernel_test(
     TestFunc,
     kernel=test_user_func_return_multiple_values,

warp/tests/test_grad_debug.py

@@ -8,7 +8,12 @@
 import unittest
 
 import warp as wp
-from warp.autograd import gradcheck, gradcheck_tape, jacobian, jacobian_fd
+from warp.autograd import (
+    gradcheck,
+    gradcheck_tape,
+    jacobian,
+    jacobian_fd,
+)
 from warp.tests.unittest_utils import *
 
 
@@ -43,7 +48,7 @@ def vec_length_kernel(a: wp.array(dtype=wp.vec3), out: wp.array(dtype=float)):
     tid = wp.tid()
     v = a[tid]
     # instead of wp.length(v), we use a trivial implementation that
-    # fails when a division by zero is occurs in the backward pass of sqrt
+    # fails when a division by zero occurs in the backward pass of sqrt
     out[tid] = wp.sqrt(v[0] ** 2.0 + v[1] ** 2.0 + v[2] ** 2.0)
 
 
@@ -63,6 +68,16 @@ def wrong_grad_kernel(a: wp.array(dtype=float), out: wp.array(dtype=float)):
     out[tid] = wrong_grad_func(a[tid])
 
 
+@wp.kernel
+def transform_point_kernel(
+    transforms: wp.array(dtype=wp.transform),
+    points: wp.array(dtype=wp.vec3),
+    out: wp.array(dtype=wp.vec3),
+):
+    tid = wp.tid()
+    out[tid] = wp.transform_point(transforms[tid], points[tid])
+
+
 def test_gradcheck_3d(test, device):
     a_3d = wp.array([((2.0, 0.0), (1.0, 0.0), (2.0, 0.0))], dtype=float, requires_grad=True, device=device)
     b_3d = wp.array([((3.0, 0.0), (1.0, 0.0), (2.0, 0.0))], dtype=float, requires_grad=True, device=device)
@@ -229,6 +244,76 @@ def test_gradcheck_tape(test, device):
     assert passed
 
 
+def test_gradcheck_function(test, device):
+    def compute_transformed_point_norms(transforms, points):
+        tf_points = wp.empty_like(points)
+        norms = wp.empty(len(points), dtype=float, requires_grad=points.requires_grad, device=points.device)
+
+        wp.launch(
+            transform_point_kernel,
+            dim=len(points),
+            inputs=[transforms, points],
+            outputs=[tf_points],
+            device=device,
+        )
+        wp.launch(
+            vec_length_kernel,
+            dim=len(points),
+            inputs=[tf_points],
+            outputs=[norms],
+            device=device,
+        )
+        return tf_points, norms
+
+    transforms = wp.array(
+        [
+            wp.transform(wp.vec3(1.0, 0.6, -2.0), wp.quat_rpy(-0.5, 0.1, 0.8)),
+            wp.transform(wp.vec3(0.2, 1.4, -0.4), wp.quat_rpy(0.5, 0.65, -0.3)),
+            wp.transform(wp.vec3(0.5, 0.2, 0.0), wp.quat_rpy(-0.5, -0.3, 0.4)),
+        ],
+        dtype=wp.transform,
+        requires_grad=True,
+        device=device,
+    )
+    points = wp.array(
+        [
+            (1.0, -0.5, 2.0),
+            (-0.95, -0.1, 0.0),
+            (9.1, 9.7, 3.8),
+        ],
+        dtype=wp.vec3,
+        requires_grad=True,
+        device=device,
+    )
+
+    jacs_ad = jacobian(
+        kernel_mixed,
+        dim=len(points),
+        inputs=[transforms, points],
+    )
+    jacs_fd = jacobian_fd(
+        kernel_mixed,
+        dim=len(points),
+        inputs=[transforms, points],
+        eps=1e-4,
+    )
+
+    # manual gradcheck
+    for i in range(2):
+        for j in range(2):
+            assert np.allclose(jacs_ad[(i, j)].numpy(), jacs_fd[(i, j)].numpy(), atol=1e-2, rtol=1e-2)
+
+    passed = gradcheck(
+        kernel_mixed,
+        dim=len(points),
+        inputs=[transforms, points],
+        raise_exception=False,
+        show_summary=False,
+    )
+
+    assert passed
+
+
 devices = get_test_devices()
 
 

warp/tests/test_hash_grid.py

@@ -85,7 +85,7 @@ def test_hashgrid_query(test, device):
 
     for i in range(num_runs):
         if print_enabled:
-            print(f"Run: {i+1}")
+            print(f"Run: {i + 1}")
             print("---------")
 
         points = particle_grid(16, 32, 16, (0.0, 0.3, 0.0), cell_radius * 0.25, 0.1)