warp-lang 1.0.0b2__py3-none-win_amd64.whl → 1.0.0b6__py3-none-win_amd64.whl

warp/tape.py

@@ -37,6 +37,7 @@ class Tape:
         print(tape.gradients[a])
 
     """
+
     def __init__(self):
         self.gradients = {}
         self.const_gradients = set()
@@ -81,7 +82,7 @@ class Tape:
             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()"
                 )
@@ -105,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 = []
@@ -129,11 +131,12 @@ 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])
 
     def record_func(self, backward, arrays):
         """

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")

warp/tests/test_adam.py

@@ -5,16 +5,14 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-# include parent path
+import unittest
+
 import numpy as np
-import math
 
 import warp as wp
-from warp.tests.test_base import *
-import unittest
-
 import warp.optim
 import warp.sim
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -28,32 +26,32 @@ def objective(params: wp.array(dtype=float), score: wp.array(dtype=float)):
 
 # This test inspired by https://machinelearningmastery.com/adam-optimization-from-scratch/
 def test_adam_solve_float(test, device):
-    wp.set_device(device)
-    params_start = np.array([0.1, 0.2], dtype=float)
-    score = wp.zeros(1, dtype=float, requires_grad=True)
-    params = wp.array(params_start, dtype=float, requires_grad=True)
-    tape = wp.Tape()
-    opt = warp.optim.Adam([params], lr=0.02, betas=(0.8, 0.999))
-
-    def gradient_func():
-        tape.reset()
-        score.zero_()
-        with tape:
-            wp.launch(kernel=objective, dim=len(params), inputs=[params, score])
-        tape.backward(score)
-        return [tape.gradients[params]]
-
-    niters = 100
-
-    opt.reset_internal_state()
-    for _ in range(niters):
-        opt.step(gradient_func())
-
-    result = params.numpy()
-    # optimum is at the origin, so the result should be close to it in all N dimensions.
-    tol = 1e-5
-    for r in result:
-        test.assertLessEqual(r, tol)
+    with wp.ScopedDevice(device):
+        params_start = np.array([0.1, 0.2], dtype=float)
+        score = wp.zeros(1, dtype=float, requires_grad=True)
+        params = wp.array(params_start, dtype=float, requires_grad=True)
+        tape = wp.Tape()
+        opt = warp.optim.Adam([params], lr=0.02, betas=(0.8, 0.999))
+
+        def gradient_func():
+            tape.reset()
+            score.zero_()
+            with tape:
+                wp.launch(kernel=objective, dim=len(params), inputs=[params, score])
+            tape.backward(score)
+            return [tape.gradients[params]]
+
+        niters = 100
+
+        opt.reset_internal_state()
+        for _ in range(niters):
+            opt.step(gradient_func())
+
+        result = params.numpy()
+        # optimum is at the origin, so the result should be close to it in all N dimensions.
+        tol = 1e-5
+        for r in result:
+            test.assertLessEqual(r, tol)
 
 
 @wp.kernel
@@ -65,32 +63,32 @@ def objective_vec3(params: wp.array(dtype=wp.vec3), score: wp.array(dtype=float)
 
 # This test inspired by https://machinelearningmastery.com/adam-optimization-from-scratch/
 def test_adam_solve_vec3(test, device):
-    wp.set_device(device)
-    params_start = np.array([[0.1, 0.2, -0.1]], dtype=float)
-    score = wp.zeros(1, dtype=float, requires_grad=True)
-    params = wp.array(params_start, dtype=wp.vec3, requires_grad=True)
-    tape = wp.Tape()
-    opt = warp.optim.Adam([params], lr=0.02, betas=(0.8, 0.999))
-
-    def gradient_func():
-        tape.reset()
-        score.zero_()
-        with tape:
-            wp.launch(kernel=objective_vec3, dim=len(params), inputs=[params, score])
-        tape.backward(score)
-        return [tape.gradients[params]]
-
-    niters = 100
-    opt.reset_internal_state()
-    for _ in range(niters):
-        opt.step(gradient_func())
-
-    result = params.numpy()
-    tol = 1e-5
-    # optimum is at the origin, so the result should be close to it in all N dimensions.
-    for r in result:
-        for v in r:
-            test.assertLessEqual(v, tol)
+    with wp.ScopedDevice(device):
+        params_start = np.array([[0.1, 0.2, -0.1]], dtype=float)
+        score = wp.zeros(1, dtype=float, requires_grad=True)
+        params = wp.array(params_start, dtype=wp.vec3, requires_grad=True)
+        tape = wp.Tape()
+        opt = warp.optim.Adam([params], lr=0.02, betas=(0.8, 0.999))
+
+        def gradient_func():
+            tape.reset()
+            score.zero_()
+            with tape:
+                wp.launch(kernel=objective_vec3, dim=len(params), inputs=[params, score])
+            tape.backward(score)
+            return [tape.gradients[params]]
+
+        niters = 100
+        opt.reset_internal_state()
+        for _ in range(niters):
+            opt.step(gradient_func())
+
+        result = params.numpy()
+        tol = 1e-5
+        # optimum is at the origin, so the result should be close to it in all N dimensions.
+        for r in result:
+            for v in r:
+                test.assertLessEqual(v, tol)
 
 
 @wp.kernel
@@ -105,56 +103,55 @@ def objective_two_inputs_vec3(
 
 # This test inspired by https://machinelearningmastery.com/adam-optimization-from-scratch/
 def test_adam_solve_two_inputs(test, device):
-    wp.set_device(device)
-    params_start1 = np.array([[0.1, 0.2, -0.1]], dtype=float)
-    params_start2 = np.array([[0.2, 0.1, 0.1]], dtype=float)
-    score = wp.zeros(1, dtype=float, requires_grad=True)
-    params1 = wp.array(params_start1, dtype=wp.vec3, requires_grad=True)
-    params2 = wp.array(params_start2, dtype=wp.vec3, requires_grad=True)
-    tape = wp.Tape()
-    opt = warp.optim.Adam([params1, params2], lr=0.02, betas=(0.8, 0.999))
-
-    def gradient_func():
-        tape.reset()
-        score.zero_()
-        with tape:
-            wp.launch(kernel=objective_two_inputs_vec3, dim=len(params1), inputs=[params1, params2, score])
-        tape.backward(score)
-        return [tape.gradients[params1], tape.gradients[params2]]
-
-    niters = 100
-    opt.reset_internal_state()
-    for _ in range(niters):
-        opt.step(gradient_func())
-
-    result = params1.numpy()
-    tol = 1e-5
-    # optimum is at the origin, so the result should be close to it in all N dimensions.
-    for r in result:
-        for v in r:
-            test.assertLessEqual(v, tol)
-
-    result = params2.numpy()
-    tol = 1e-5
-    # optimum is at the origin, so the result should be close to it in all N dimensions.
-    for r in result:
-        for v in r:
-            test.assertLessEqual(v, tol)
-
-
-def register(parent):
-    devices = get_test_devices()
-
-    class TestArray(parent):
-        pass
-
-    add_function_test(TestArray, "test_adam_solve_float", test_adam_solve_float, devices=devices)
-    add_function_test(TestArray, "test_adam_solve_vec3", test_adam_solve_vec3, devices=devices)
-    add_function_test(TestArray, "test_adam_solve_two_inputs", test_adam_solve_two_inputs, devices=devices)
-
-    return TestArray
+    with wp.ScopedDevice(device):
+        params_start1 = np.array([[0.1, 0.2, -0.1]], dtype=float)
+        params_start2 = np.array([[0.2, 0.1, 0.1]], dtype=float)
+        score = wp.zeros(1, dtype=float, requires_grad=True)
+        params1 = wp.array(params_start1, dtype=wp.vec3, requires_grad=True)
+        params2 = wp.array(params_start2, dtype=wp.vec3, requires_grad=True)
+        tape = wp.Tape()
+        opt = warp.optim.Adam([params1, params2], lr=0.02, betas=(0.8, 0.999))
+
+        def gradient_func():
+            tape.reset()
+            score.zero_()
+            with tape:
+                wp.launch(kernel=objective_two_inputs_vec3, dim=len(params1), inputs=[params1, params2, score])
+            tape.backward(score)
+            return [tape.gradients[params1], tape.gradients[params2]]
+
+        niters = 100
+        opt.reset_internal_state()
+        for _ in range(niters):
+            opt.step(gradient_func())
+
+        result = params1.numpy()
+        tol = 1e-5
+        # optimum is at the origin, so the result should be close to it in all N dimensions.
+        for r in result:
+            for v in r:
+                test.assertLessEqual(v, tol)
+
+        result = params2.numpy()
+        tol = 1e-5
+        # optimum is at the origin, so the result should be close to it in all N dimensions.
+        for r in result:
+            for v in r:
+                test.assertLessEqual(v, tol)
+
+
+devices = get_test_devices()
+
+
+class TestAdam(unittest.TestCase):
+    pass
+
+
+add_function_test(TestAdam, "test_adam_solve_float", test_adam_solve_float, devices=devices)
+add_function_test(TestAdam, "test_adam_solve_vec3", test_adam_solve_vec3, devices=devices)
+add_function_test(TestAdam, "test_adam_solve_two_inputs", test_adam_solve_two_inputs, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)