warp-lang 0.10.1__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/tape.py

@@ -5,11 +5,39 @@
 # 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
 
 
 class Tape:
+    """
+    Record kernel launches within a Tape scope to enable automatic differentiation.
+    Gradients can be computed after the operations have been recorded on the tape via
+    ``tape.backward()``.
+
+    Example
+    -------
+
+    .. code-block:: python
+
+        tape = wp.Tape()
+
+        # forward pass
+        with tape:
+            wp.launch(kernel=compute1, inputs=[a, b], device="cuda")
+            wp.launch(kernel=compute2, inputs=[c, d], device="cuda")
+            wp.launch(kernel=loss, inputs=[d, l], device="cuda")
+
+        # reverse pass
+        tape.backward(l)
+
+    Gradients can be accessed via the ``tape.gradients`` dictionary, e.g.:
+
+    .. code-block:: python
+
+        print(tape.gradients[a])
+
+    """
+
     def __init__(self):
         self.gradients = {}
         self.const_gradients = set()
@@ -37,13 +65,24 @@ class Tape:
     #  adj_tensor = tape.gradients[tensor]
     #
     def backward(self, loss: wp.array = None, grads: dict = None):
+        """
+        Evaluate the backward pass of the recorded operations on the tape.
+        A single-element array ``loss`` or a dictionary of arrays ``grads``
+        can be provided to assign the incoming gradients for the reverse-mode
+        automatic differentiation pass.
+
+        Args:
+            loss (wp.array): A single-element array that holds the loss function value whose gradient is to be computed
+            grads (dict): A dictionary of arrays that map from Warp arrays to their incoming gradients
+
+        """
         # if scalar loss is specified then initialize
         # a 'seed' array for it, with gradient of one
         if loss:
             if loss.size > 1 or wp.types.type_length(loss.dtype) > 1:
                 raise RuntimeError("Can only return gradients for scalar loss functions.")
 
-            if loss.requires_grad == False:
+            if not loss.requires_grad:
                 raise RuntimeError(
                     "Scalar loss arrays should have requires_grad=True set before calling Tape.backward()"
                 )
@@ -67,9 +106,10 @@ class Tape:
             else:
                 kernel = launch[0]
                 dim = launch[1]
-                inputs = launch[2]
-                outputs = launch[3]
-                device = launch[4]
+                max_blocks = launch[2]
+                inputs = launch[3]
+                outputs = launch[4]
+                device = launch[5]
 
                 adj_inputs = []
                 adj_outputs = []
@@ -91,16 +131,21 @@ class Tape:
                     adj_outputs=adj_outputs,
                     device=device,
                     adjoint=True,
+                    max_blocks=max_blocks,
                 )
 
     # record a kernel launch on the tape
-    def record_launch(self, kernel, dim, inputs, outputs, device):
-        self.launches.append([kernel, dim, inputs, outputs, device])
+    def record_launch(self, kernel, dim, max_blocks, inputs, outputs, device):
+        self.launches.append([kernel, dim, max_blocks, inputs, outputs, device])
 
-    # records a custom function for the backward pass, can be any
-    # Callable python object. Callee should also pass arrays that
-    # take part in the function for gradient tracking.
     def record_func(self, backward, arrays):
+        """
+        Records a custom function to be executed only in the backward pass.
+
+        Args:
+            backward (Callable): A callable Python object (can be any function) that will be executed in the backward pass.
+            arrays (list): A list of arrays that are used by the function for gradient tracking.
+        """
         self.launches.append(backward)
 
         for a in arrays:
@@ -145,10 +190,16 @@ class Tape:
         return None
 
     def reset(self):
+        """
+        Clear all operations recorded on the tape and zero out all gradients.
+        """
         self.launches = []
         self.zero()
 
     def zero(self):
+        """
+        Zero out all gradients recorded on the tape.
+        """
         for a, g in self.gradients.items():
             if a not in self.const_gradients:
                 if isinstance(a, wp.codegen.StructInstance):

warp/tests/__main__.py

@@ -1,7 +1,4 @@
-from .test_all import run
+from warp.thirdparty.unittest_parallel import main
 
-ret = run()
-
-import sys
-
-sys.exit(ret)
+if __name__ == "__main__":
+    main()

warp/tests/{test_class_kernel.py → aux_test_class_kernel.py}

@@ -1,7 +1,15 @@
+# Copyright (c) 2022 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.
+
+"""Dummy class used in test_reload.py"""
+
 import warp as wp
 
 
-# dummy class used in test_reload.py
 class ClassKernelTest:
     def __init__(self, device):
         # 3x3 frames in the rest pose:

warp/tests/aux_test_conditional_unequal_types_kernels.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2023 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.
+
+"""This file defines a kernel that fails on codegen.py"""
+
+import warp as wp
+
+
+@wp.kernel
+def unequal_types_kernel():
+    x = wp.int32(10)
+    y = 10
+    z = True
+
+    # Throws a TypeError
+    if x == y == z:
+        pass

warp/tests/{test_dependent.py → aux_test_dependent.py}

@@ -5,10 +5,10 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-# This file is used to test reloading module references.
+"""This file is used to test reloading module references."""
 
 import warp as wp
-import warp.tests.test_reference as ref
+import warp.tests.aux_test_reference as ref
 
 wp.init()
 

warp/tests/{test_reference.py → aux_test_reference.py}

@@ -1,7 +1,7 @@
 # This file is used to test reloading module references.
 
 import warp as wp
-import warp.tests.test_reference_reference as refref
+import warp.tests.aux_test_reference_reference as refref
 
 wp.init()
 

warp/tests/aux_test_unresolved_func.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2022 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 warp as wp
+
+
+@wp.kernel
+def unresolved_func_kernel():
+    # this should trigger an exception due to unresolved function
+    x = wp.missing_func(42)

warp/tests/aux_test_unresolved_symbol.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2022 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 warp as wp
+
+
+@wp.kernel
+def unresolved_symbol_kernel():
+    # this should trigger an exception due to unresolved symbol
+    x = missing_symbol

warp/tests/disabled_kinematics.py

@@ -0,0 +1,239 @@
+# Copyright (c) 2022 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 math
+import unittest
+
+import warp as wp
+import warp.sim
+from warp.tests.unittest_utils import *
+
+wp.init()
+
+
+def build_ant(num_envs):
+    builder = wp.sim.ModelBuilder()
+    for i in range(num_envs):
+        wp.sim.parse_mjcf(
+            os.path.join(os.path.dirname(__file__), "../../examples/assets/nv_ant.xml"),
+            builder,
+            up_axis="y",
+        )
+
+        coord_count = 15
+        dof_count = 14
+
+        coord_start = i * coord_count
+        dof_start = i * dof_count
+
+        # base
+        builder.joint_q[coord_start : coord_start + 3] = [i * 2.0, 0.70, 0.0]
+        builder.joint_q[coord_start + 3 : coord_start + 7] = wp.quat_from_axis_angle((1.0, 0.0, 0.0), -math.pi * 0.5)
+
+        # joints
+        builder.joint_q[coord_start + 7 : coord_start + coord_count] = [0.0, 1.0, 0.0, -1.0, 0.0, -1.0, 0.0, 1.0]
+        builder.joint_qd[dof_start + 6 : dof_start + dof_count] = [1.0, 1.0, 1.0, -1.0, 1.0, -1.0, 1.0, 1.0]
+
+    return builder
+
+
+def build_complex_joint_mechanism(chain_length):
+    builder = wp.sim.ModelBuilder()
+    com0 = wp.vec3(1.0, 2.0, 3.0)
+    com1 = wp.vec3(4.0, 5.0, 6.0)
+    com2 = wp.vec3(7.0, 8.0, 9.0)
+    ax0 = wp.normalize(wp.vec3(-1.0, 2.0, 3.0))
+    ax1 = wp.normalize(wp.vec3(4.0, -1.0, 2.0))
+    ax2 = wp.normalize(wp.vec3(-3.0, 4.0, -1.0))
+    # declare some transforms with nonzero translation and orientation
+    tf0 = wp.transform(wp.vec3(1.0, 2.0, 3.0), wp.quat_from_axis_angle((1.0, 0.0, 0.0), math.pi * 0.25))
+    tf1 = wp.transform(wp.vec3(4.0, 5.0, 6.0), wp.quat_from_axis_angle((0.0, 1.0, 0.0), math.pi * 0.5))
+    tf2 = wp.transform(wp.vec3(7.0, 8.0, 9.0), wp.quat_from_axis_angle((0.0, 0.0, 1.0), math.pi * 0.75))
+
+    parent = -1
+    for i in range(chain_length):
+        b0 = builder.add_body(com=com0)
+        builder.add_joint_fixed(parent=parent, child=b0, parent_xform=tf1, child_xform=tf0)
+        assert builder.articulation_count == 1
+
+        b1 = builder.add_body(com=com1)
+        builder.add_joint_revolute(parent=b0, child=b1, parent_xform=tf1, child_xform=tf2, axis=ax1)
+        builder.joint_q[-1] = 0.3
+        builder.joint_qd[-1] = 1.0
+
+        b2 = builder.add_body(com=com2)
+        builder.add_joint_universal(parent=b1, child=b2, parent_xform=tf2, child_xform=tf0, axis_0=ax0, axis_1=ax1)
+        builder.joint_q[-2:] = [0.3, 0.5]
+        builder.joint_qd[-2:] = [1.0, -1.0]
+
+        b3 = builder.add_body(com=com0)
+        builder.add_joint_ball(parent=b2, child=b3, parent_xform=tf0, child_xform=tf1)
+        builder.joint_q[-4:] = list(wp.quat_from_axis_angle(ax0, 0.7))
+        builder.joint_qd[-3:] = [1.0, -0.6, 1.5]
+
+        b4 = builder.add_body(com=com1)
+        builder.add_joint_compound(
+            parent=b3,
+            child=b4,
+            parent_xform=tf2,
+            child_xform=tf1,
+            axis_0=(0, 0, 1),
+            axis_1=(1, 0, 0),
+            axis_2=(0, 1, 0),
+        )
+        builder.joint_q[-3:] = [0.3, 0.5, 0.27]
+        builder.joint_qd[-3:] = [1.23, -1.0, 0.5]
+
+        b5 = builder.add_body(com=com2)
+        builder.add_joint_prismatic(
+            parent=b4,
+            child=b5,
+            parent_xform=tf2,
+            child_xform=tf0,
+            axis=ax0,
+        )
+        builder.joint_q[-1] = 0.92
+        builder.joint_qd[-1] = -0.63
+
+        b6 = builder.add_body(com=com0)
+        builder.add_joint_d6(
+            parent=b5,
+            child=b6,
+            parent_xform=tf0,
+            child_xform=tf2,
+            linear_axes=[ax0, ax1, wp.cross(ax0, ax1)],
+            angular_axes=[ax1, ax2, wp.cross(ax1, ax2)],
+        )
+        builder.joint_q[-6:] = [0.3, 0.5, 0.7, 0.9, 1.1, 1.3]
+        builder.joint_qd[-6:] = [1.0, -1.0, 0.5, 0.8, -0.3, 0.1]
+
+        b7 = builder.add_body(com=com1)
+        builder.add_joint_free(
+            parent=b6,
+            child=b7,
+            parent_xform=tf1,
+            child_xform=tf2,
+        )
+        builder.joint_q[-7:] = [0.5, -0.9, 1.4] + list(wp.quat_rpy(0.3, -0.5, 0.7))
+        builder.joint_qd[-6:] = [1.0, -1.0, 0.5, 0.8, -0.3, 0.1]
+
+        b8 = builder.add_body(com=com2)
+        builder.add_joint_distance(
+            parent=b7,
+            child=b8,
+            parent_xform=tf1,
+            child_xform=tf2,
+        )
+        builder.joint_q[-7:] = [-0.3, -0.7, 0.2] + list(wp.quat_rpy(0.1, 0.1, 0.4))
+        builder.joint_qd[-6:] = [-0.34, 0.5, -0.6, -0.4, 0.2, 0.1]
+
+        # D6 joint that behaves like a fixed joint
+        b9 = builder.add_body(com=com0)
+        builder.add_joint_d6(
+            parent=b8,
+            child=b9,
+            parent_xform=tf0,
+            child_xform=tf2,
+            linear_axes=[],
+            angular_axes=[],
+        )
+
+        b10 = builder.add_body(com=com0)
+        builder.add_joint_d6(
+            parent=b9,
+            child=b10,
+            parent_xform=tf1,
+            child_xform=tf2,
+            linear_axes=[ax1],
+            angular_axes=[ax2, ax0],
+        )
+        builder.joint_q[-3:] = [0.3, 0.5, 0.7]
+        builder.joint_qd[-3:] = [1.0, -1.0, 0.5]
+
+        b11 = builder.add_body(com=com1)
+        builder.add_joint_d6(
+            parent=b10,
+            child=b11,
+            parent_xform=tf1,
+            child_xform=tf2,
+            linear_axes=[ax1, ax0, wp.cross(ax1, ax0)],
+            angular_axes=[],
+        )
+        builder.joint_q[-3:] = [0.3, 0.5, 0.7]
+        builder.joint_qd[-3:] = [1.0, -1.0, 0.5]
+
+        b12 = builder.add_body(com=com2)
+        builder.add_joint_d6(
+            parent=b11,
+            child=b12,
+            parent_xform=tf1,
+            child_xform=tf2,
+            linear_axes=[],
+            angular_axes=[ax1, ax2, wp.cross(ax1, ax2)],
+        )
+        builder.joint_q[-3:] = [0.3, 0.5, 0.7]
+        builder.joint_qd[-3:] = [1.0, -1.0, 0.5]
+
+        parent = b12
+
+    return builder
+
+
+def check_fk_ik(builder, device):
+    model = builder.finalize(device)
+    state = model.state()
+
+    q_fk = model.joint_q.numpy()
+    qd_fk = model.joint_qd.numpy()
+
+    wp.sim.eval_fk(model, model.joint_q, model.joint_qd, None, state)
+
+    q_ik = wp.zeros_like(model.joint_q)
+    qd_ik = wp.zeros_like(model.joint_qd)
+
+    wp.sim.eval_ik(model, state, q_ik, qd_ik)
+
+    # adjust numpy print settings
+    # np.set_printoptions(precision=4, floatmode="fixed", suppress=True)
+    # print("q:")
+    # print(np.array(q_fk))
+    # print(q_ik.numpy())
+
+    # print("qd:")
+    # print(np.array(qd_fk))
+    # print(qd_ik.numpy())
+
+    assert_np_equal(q_ik.numpy(), q_fk, tol=1e-4)
+    assert_np_equal(qd_ik.numpy(), qd_fk, tol=1e-4)
+
+
+def test_fk_ik_ant(test, device):
+    builder = build_ant(3)
+    check_fk_ik(builder, device)
+
+
+def test_fk_ik_complex_joint_mechanism(test, device):
+    builder = build_complex_joint_mechanism(2)
+    check_fk_ik(builder, device)
+
+
+devices = get_test_devices()
+
+
+class TestKinematics(unittest.TestCase):
+    pass
+
+
+add_function_test(TestKinematics, "test_fk_ik_ant", test_fk_ik_ant, devices=devices)
+add_function_test(
+    TestKinematics, "test_fk_ik_complex_joint_mechanism", test_fk_ik_complex_joint_mechanism, devices=devices
+)
+
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=False)

warp/tests/run_coverage_serial.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2023 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.
+
+"""Serial code-coverage runner
+
+This script is used to generate code-coverage reports by running Warp tests.
+It runs in serial so can take over an hour to finish. To generate a coverage
+report in parallel, use the warp/thirdparty./unittest_parallel.py script
+instead with the --coverage option, e.g. python -m warp.tests --coverage
+"""
+
+import coverage
+
+cover = coverage.Coverage(config_file=True, messages=True)
+
+cover.start()
+
+with cover.collect():
+    import unittest_serial  # noqa: E402
+
+    unittest_serial.run_specified()
+
+cover.save()
+
+cover.report()
+
+cover.html_report(title="Warp Testing Code Coverage Report")