warp-lang 1.0.0b5__py3-none-manylinux2014_x86_64.whl → 1.0.0b6__py3-none-manylinux2014_x86_64.whl

warp/stubs.py

@@ -40,7 +40,8 @@ from warp.types import spatial_vector, spatial_vectorh, spatial_vectorf, spatial
 from warp.types import spatial_matrix, spatial_matrixh, spatial_matrixf, spatial_matrixd
 
 from warp.types import Bvh, Mesh, HashGrid, Volume, MarchingCubes
-from warp.types import bvh_query_t, mesh_query_aabb_t, hash_grid_query_t
+from warp.types import bvh_query_t, hash_grid_query_t, mesh_query_aabb_t, mesh_query_point_t, mesh_query_ray_t
+
 
 from warp.types import matmul, adj_matmul, batched_matmul, adj_batched_matmul, from_ptr
 
@@ -69,7 +70,6 @@ from warp.context import (
 )
 from warp.context import set_module_options, get_module_options, get_module
 from warp.context import capture_begin, capture_end, capture_launch
-from warp.context import print_builtins, export_builtins, export_stubs
 from warp.context import Kernel, Function, Launch
 from warp.context import Stream, get_stream, set_stream, synchronize_stream
 from warp.context import Event, record_event, wait_event, wait_stream
@@ -868,7 +868,7 @@ def mlp(
 def bvh_query_aabb(id: uint64, lower: vec3f, upper: vec3f) -> bvh_query_t:
     """
     Construct an axis-aligned bounding box query against a BVH object. This query can be used to iterate over all bounds
-       inside a BVH. Returns an object that is used to track state during BVH traversal.
+       inside a BVH.
 
        :param id: The BVH identifier
        :param lower: The lower bound of the bounding box in BVH space
@@ -881,7 +881,7 @@ def bvh_query_aabb(id: uint64, lower: vec3f, upper: vec3f) -> bvh_query_t:
 def bvh_query_ray(id: uint64, start: vec3f, dir: vec3f) -> bvh_query_t:
     """
     Construct a ray query against a BVH object. This query can be used to iterate over all bounds
-       that intersect the ray. Returns an object that is used to track state during BVH traversal.
+       that intersect the ray.
 
        :param id: The BVH identifier
        :param start: The start of the ray in BVH space
@@ -900,11 +900,9 @@ def bvh_query_next(query: bvh_query_t, index: int32) -> bool:
 
 
 @over
-def mesh_query_point(
-    id: uint64, point: vec3f, max_dist: float32, inside: float32, face: int32, bary_u: float32, bary_v: float32
-) -> bool:
+def mesh_query_point(id: uint64, point: vec3f, max_dist: float32) -> mesh_query_point_t:
     """
-    Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given ``point`` in space. Returns ``True`` if a point < ``max_dist`` is found.
+    Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given ``point`` in space.
 
        Identifies the sign of the distance using additional ray-casts to determine if the point is inside or outside.
        This method is relatively robust, but does increase computational cost.
@@ -913,66 +911,42 @@ def mesh_query_point(
        :param id: The mesh identifier
        :param point: The point in space to query
        :param max_dist: Mesh faces above this distance will not be considered by the query
-       :param inside: Returns a value < 0 if query point is inside the mesh, >=0 otherwise.
-                      Note that mesh must be watertight for this to be robust
-       :param face: Returns the index of the closest face
-       :param bary_u: Returns the barycentric u coordinate of the closest point
-       :param bary_v: Returns the barycentric v coordinate of the closest point
     """
     ...
 
 
 @over
-def mesh_query_point_no_sign(
-    id: uint64, point: vec3f, max_dist: float32, face: int32, bary_u: float32, bary_v: float32
-) -> bool:
+def mesh_query_point_no_sign(id: uint64, point: vec3f, max_dist: float32) -> mesh_query_point_t:
     """
-    Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given ``point`` in space. Returns ``True`` if a point < ``max_dist`` is found.
+    Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given ``point`` in space.
 
        This method does not compute the sign of the point (inside/outside) which makes it faster than other point query methods.
 
        :param id: The mesh identifier
        :param point: The point in space to query
        :param max_dist: Mesh faces above this distance will not be considered by the query
-       :param face: Returns the index of the closest face
-       :param bary_u: Returns the barycentric u coordinate of the closest point
-       :param bary_v: Returns the barycentric v coordinate of the closest point
     """
     ...
 
 
 @over
-def mesh_query_furthest_point_no_sign(
-    id: uint64, point: vec3f, min_dist: float32, face: int32, bary_u: float32, bary_v: float32
-) -> bool:
+def mesh_query_furthest_point_no_sign(id: uint64, point: vec3f, min_dist: float32) -> mesh_query_point_t:
     """
-    Computes the furthest point on the mesh with identifier `id` to the given point in space. Returns ``True`` if a point > ``min_dist`` is found.
+    Computes the furthest point on the mesh with identifier `id` to the given point in space.
 
        This method does not compute the sign of the point (inside/outside).
 
        :param id: The mesh identifier
        :param point: The point in space to query
        :param min_dist: Mesh faces below this distance will not be considered by the query
-       :param face: Returns the index of the furthest face
-       :param bary_u: Returns the barycentric u coordinate of the furthest point
-       :param bary_v: Returns the barycentric v coordinate of the furthest point
     """
     ...
 
 
 @over
-def mesh_query_point_sign_normal(
-    id: uint64,
-    point: vec3f,
-    max_dist: float32,
-    inside: float32,
-    face: int32,
-    bary_u: float32,
-    bary_v: float32,
-    epsilon: float32,
-) -> bool:
+def mesh_query_point_sign_normal(id: uint64, point: vec3f, max_dist: float32, epsilon: float32) -> mesh_query_point_t:
     """
-    Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given ``point`` in space. Returns ``True`` if a point < ``max_dist`` is found.
+    Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given ``point`` in space.
 
        Identifies the sign of the distance (inside/outside) using the angle-weighted pseudo normal.
        This approach to sign determination is robust for well conditioned meshes that are watertight and non-self intersecting.
@@ -981,11 +955,6 @@ def mesh_query_point_sign_normal(
        :param id: The mesh identifier
        :param point: The point in space to query
        :param max_dist: Mesh faces above this distance will not be considered by the query
-       :param inside: Returns a value < 0 if query point is inside the mesh, >=0 otherwise.
-                      Note that mesh must be watertight for this to be robust
-       :param face: Returns the index of the closest face
-       :param bary_u: Returns the barycentric u coordinate of the closest point
-       :param bary_v: Returns the barycentric v coordinate of the closest point
        :param epsilon: Epsilon treating distance values as equal, when locating the minimum distance vertex/face/edge, as a
                        fraction of the average edge length, also for treating closest point as being on edge/vertex default 1e-3
     """
@@ -994,18 +963,10 @@ def mesh_query_point_sign_normal(
 
 @over
 def mesh_query_point_sign_winding_number(
-    id: uint64,
-    point: vec3f,
-    max_dist: float32,
-    inside: float32,
-    face: int32,
-    bary_u: float32,
-    bary_v: float32,
-    accuracy: float32,
-    threshold: float32,
-) -> bool:
-    """
-    Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given point in space. Returns ``True`` if a point < ``max_dist`` is found.
+    id: uint64, point: vec3f, max_dist: float32, accuracy: float32, threshold: float32
+) -> mesh_query_point_t:
+    """
+    Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given point in space.
 
        Identifies the sign using the winding number of the mesh relative to the query point. This method of sign determination is robust for poorly conditioned meshes
        and provides a smooth approximation to sign even when the mesh is not watertight. This method is the most robust and accurate of the sign determination meshes
@@ -1016,11 +977,6 @@ def mesh_query_point_sign_winding_number(
        :param id: The mesh identifier
        :param point: The point in space to query
        :param max_dist: Mesh faces above this distance will not be considered by the query
-       :param inside: Returns a value < 0 if query point is inside the mesh, >=0 otherwise.
-                      Note that mesh must be watertight for this to be robust
-       :param face: Returns the index of the closest face
-       :param bary_u: Returns the barycentric u coordinate of the closest point
-       :param bary_v: Returns the barycentric v coordinate of the closest point
        :param accuracy: Accuracy for computing the winding number with fast winding number method utilizing second-order dipole approximation, default 2.0
        :param threshold: The threshold of the winding number to be considered inside, default 0.5
     """
@@ -1028,31 +984,14 @@ def mesh_query_point_sign_winding_number(
 
 
 @over
-def mesh_query_ray(
-    id: uint64,
-    start: vec3f,
-    dir: vec3f,
-    max_t: float32,
-    t: float32,
-    bary_u: float32,
-    bary_v: float32,
-    sign: float32,
-    normal: vec3f,
-    face: int32,
-) -> bool:
+def mesh_query_ray(id: uint64, start: vec3f, dir: vec3f, max_t: float32) -> mesh_query_ray_t:
     """
-    Computes the closest ray hit on the :class:`Mesh` with identifier ``id``, returns ``True`` if a point < ``max_t`` is found.
+    Computes the closest ray hit on the :class:`Mesh` with identifier ``id``.
 
        :param id: The mesh identifier
        :param start: The start point of the ray
        :param dir: The ray direction (should be normalized)
        :param max_t: The maximum distance along the ray to check for intersections
-       :param t: Returns the distance of the closest hit along the ray
-       :param bary_u: Returns the barycentric u coordinate of the closest hit
-       :param bary_v: Returns the barycentric v coordinate of the closest hit
-       :param sign: Returns a value > 0 if the hit ray hit front of the face, returns < 0 otherwise
-       :param normal: Returns the face normal
-       :param face: Returns the index of the hit face
     """
     ...
 
@@ -1062,7 +1001,6 @@ def mesh_query_aabb(id: uint64, lower: vec3f, upper: vec3f) -> mesh_query_aabb_t
     """
     Construct an axis-aligned bounding box query against a :class:`Mesh`.
        This query can be used to iterate over all triangles inside a volume.
-       Returns an object that is used to track state during mesh traversal.
 
        :param id: The mesh identifier
        :param lower: The lower bound of the bounding box in mesh space
@@ -1099,7 +1037,7 @@ def mesh_eval_velocity(id: uint64, face: int32, bary_u: float32, bary_v: float32
 @over
 def hash_grid_query(id: uint64, point: vec3f, max_dist: float32) -> hash_grid_query_t:
     """
-    Construct a point query against a :class:`HashGrid`. This query can be used to iterate over all neighboring points within a fixed radius from the query point. Returns an object that is used to track state during neighbor traversal.
+    Construct a point query against a :class:`HashGrid`. This query can be used to iterate over all neighboring points within a fixed radius from the query point.
     """
     ...
 

warp/tape.py

@@ -82,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 is False:
+            if not loss.requires_grad:
                 raise RuntimeError(
                     "Scalar loss arrays should have requires_grad=True set before calling Tape.backward()"
                 )

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/{test_kinematics.py → disabled_kinematics.py}

@@ -5,12 +5,12 @@
 # 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
-import math
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -221,21 +221,19 @@ def test_fk_ik_complex_joint_mechanism(test, device):
     check_fk_ik(builder, device)
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestGrad(parent):
-        pass
+class TestKinematics(unittest.TestCase):
+    pass
 
-    add_function_test(TestGrad, "test_fk_ik_ant", test_fk_ik_ant, devices=devices)
-    add_function_test(
-        TestGrad, "test_fk_ik_complex_joint_mechanism", test_fk_ik_complex_joint_mechanism, devices=devices
-    )
 
-    return TestGrad
+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()
-    _ = register(unittest.TestCase)
     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,57 +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__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)