warp-lang 1.3.3__py3-none-macosx_10_13_universal2.whl → 1.4.0__py3-none-macosx_10_13_universal2.whl

warp/tests/test_triangle_closest_point.py

@@ -0,0 +1,137 @@
+# 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 warp.sim.collide import triangle_closest_point_barycentric
+from warp.tests.unittest_utils import *
+
+
+# a-b is the edge where the closest point is located at
+@wp.func
+def check_edge_feasible_region(p: wp.vec3, a: wp.vec3, b: wp.vec3, c: wp.vec3, eps: float):
+    ap = p - a
+    bp = p - b
+    ab = b - a
+
+    if wp.dot(ap, ab) < -eps:
+        return False
+
+    if wp.dot(bp, ab) > eps:
+        return False
+
+    ab_sqr_norm = wp.dot(ab, ab)
+    if ab_sqr_norm < eps:
+        return False
+
+    t = wp.dot(ab, c - a) / ab_sqr_norm
+
+    perpendicular_foot = a + t * ab
+
+    if wp.dot(c - perpendicular_foot, p - perpendicular_foot) > eps:
+        return False
+
+    return True
+
+
+# closest point is a
+@wp.func
+def check_vertex_feasible_region(p: wp.vec3, a: wp.vec3, b: wp.vec3, c: wp.vec3, eps: float):
+    ap = p - a
+    ba = a - b
+    ca = a - c
+
+    if wp.dot(ap, ba) < -eps:
+        return False
+
+    if wp.dot(p, ca) < -eps:
+        return False
+
+    return True
+
+
+@wp.kernel
+def test_triangle_closest_point_kernel(tri: wp.array(dtype=wp.vec3), passed: wp.array(dtype=wp.bool)):
+    state = wp.uint32(wp.rand_init(wp.int32(123), wp.int32(0)))
+    eps = 1e-5
+
+    a = tri[0]
+    b = tri[1]
+    c = tri[2]
+
+    for _i in range(1000):
+        l = wp.float32(0.0)
+        while l < eps:
+            p = wp.vec3(wp.randn(state), wp.randn(state), wp.randn(state))
+            l = wp.length(p)
+
+        # project to a sphere with r=2
+        p = 2.0 * p / l
+
+        bary = triangle_closest_point_barycentric(tri[0], tri[1], tri[2], p)
+
+        for dim in range(3):
+            v1_index = (dim + 1) % 3
+            v2_index = (dim + 2) % 3
+            v1 = tri[v1_index]
+            v2 = tri[v2_index]
+            v3 = tri[dim]
+
+            # on edge
+            if bary[dim] == 0.0 and bary[v1_index] != 0.0 and bary[v2_index] != 0.0:
+                if not check_edge_feasible_region(p, v1, v2, v3, eps):
+                    passed[0] = False
+                    return
+
+                # p-closest_p must be perpendicular to v1-v2
+                closest_p = a * bary[0] + b * bary[1] + c * bary[2]
+                e = v1 - v2
+                err = wp.dot(e, closest_p - p)
+                if wp.abs(err) > eps:
+                    passed[0] = False
+                    return
+
+            if bary[v1_index] == 0.0 and bary[v2_index] == 0.0:
+                if not check_vertex_feasible_region(p, v3, v1, v2, eps):
+                    passed[0] = False
+                    return
+
+            if bary[dim] != 0.0 and bary[v1_index] != 0.0 and bary[v2_index] != 0.0:
+                closest_p = a * bary[0] + b * bary[1] + c * bary[2]
+                e1 = v1 - v2
+                e2 = v1 - v3
+                if wp.abs(wp.dot(e1, closest_p - p)) > eps or wp.abs(wp.dot(e2, closest_p - p)) > eps:
+                    passed[0] = False
+                    return
+
+
+def test_triangle_closest_point(test, device):
+    passed = wp.array([True], dtype=wp.bool, device=device)
+
+    a = wp.vec3(1.0, 0.0, 0.0)
+    b = wp.vec3(0.0, 0.0, 0.0)
+    c = wp.vec3(0.0, 1.0, 0.0)
+
+    tri = wp.array([a, b, c], dtype=wp.vec3, device=device)
+    wp.launch(test_triangle_closest_point_kernel, dim=1, inputs=[tri, passed], device=device)
+    passed = passed.numpy()
+
+    test.assertTrue(passed.all())
+
+
+devices = get_test_devices()
+
+
+class TestTriangleClosestPoint(unittest.TestCase):
+    pass
+
+
+add_function_test(TestTriangleClosestPoint, "test_triangle_closest_point", test_triangle_closest_point, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_types.py

@@ -215,7 +215,7 @@ class TestTypes(unittest.TestCase):
         self.assertEqual(const, wp.vec3i(1, 2, 3))
 
     def test_constant_error_invalid_type(self):
-        with self.assertRaisesRegex(RuntimeError, r"Invalid constant type: <class 'tuple'>$"):
+        with self.assertRaisesRegex(TypeError, r"Invalid constant type: <class 'tuple'>$"):
             wp.constant((1, 2, 3))
 
     def test_vector_assign(self):

warp/tests/test_vbd.py

@@ -0,0 +1,386 @@
+# 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
+import warp.optim
+import warp.sim
+from warp.sim.model import (
+    PARTICLE_FLAG_ACTIVE,
+)
+from warp.tests.unittest_utils import *
+
+
+class VBDClothSim:
+    def __init__(self, device):
+        # fmt: off
+        self.pts = [
+            (-50.0000000, 0.0000000, -50.0000000),
+            (-38.8888893, 11.1111107, -50.0000000),
+            (-27.7777786, 22.2222214, -50.0000000),
+            (-16.6666679, 33.3333321, -50.0000000),
+            (-5.5555558, 44.4444427, -50.0000000),
+            (5.5555558, 55.5555573, -50.0000000),
+            (16.6666679, 66.6666641, -50.0000000),
+            (27.7777786, 77.7777786, -50.0000000),
+            (38.8888893, 88.8888855, -50.0000000),
+            (50.0000000, 100.0000000, -50.0000000),
+            (-50.0000000, 0.0000000, -38.8888893),
+            (-38.8888893, 11.1111107, -38.8888893),
+            (-27.7777786, 22.2222214, -38.8888893),
+            (-16.6666679, 33.3333321, -38.8888893),
+            (-5.5555558, 44.4444427, -38.8888893),
+            (5.5555558, 55.5555573, -38.8888893),
+            (16.6666679, 66.6666641, -38.8888893),
+            (27.7777786, 77.7777786, -38.8888893),
+            (38.8888893, 88.8888855, -38.8888893),
+            (50.0000000, 100.0000000, -38.8888893),
+            (-50.0000000, 0.0000000, -27.7777786),
+            (-38.8888893, 11.1111107, -27.7777786),
+            (-27.7777786, 22.2222214, -27.7777786),
+            (-16.6666679, 33.3333321, -27.7777786),
+            (-5.5555558, 44.4444427, -27.7777786),
+            (5.5555558, 55.5555573, -27.7777786),
+            (16.6666679, 66.6666641, -27.7777786),
+            (27.7777786, 77.7777786, -27.7777786),
+            (38.8888893, 88.8888855, -27.7777786),
+            (50.0000000, 100.0000000, -27.7777786),
+            (-50.0000000, 0.0000000, -16.6666679),
+            (-38.8888893, 11.1111107, -16.6666679),
+            (-27.7777786, 22.2222214, -16.6666679),
+            (-16.6666679, 33.3333321, -16.6666679),
+            (-5.5555558, 44.4444427, -16.6666679),
+            (5.5555558, 55.5555573, -16.6666679),
+            (16.6666679, 66.6666641, -16.6666679),
+            (27.7777786, 77.7777786, -16.6666679),
+            (38.8888893, 88.8888855, -16.6666679),
+            (50.0000000, 100.0000000, -16.6666679),
+            (-50.0000000, 0.0000000, -5.5555558),
+            (-38.8888893, 11.1111107, -5.5555558),
+            (-27.7777786, 22.2222214, -5.5555558),
+            (-16.6666679, 33.3333321, -5.5555558),
+            (-5.5555558, 44.4444427, -5.5555558),
+            (5.5555558, 55.5555573, -5.5555558),
+            (16.6666679, 66.6666641, -5.5555558),
+            (27.7777786, 77.7777786, -5.5555558),
+            (38.8888893, 88.8888855, -5.5555558),
+            (50.0000000, 100.0000000, -5.5555558),
+            (-50.0000000, 0.0000000, 5.5555558),
+            (-38.8888893, 11.1111107, 5.5555558),
+            (-27.7777786, 22.2222214, 5.5555558),
+            (-16.6666679, 33.3333321, 5.5555558),
+            (-5.5555558, 44.4444427, 5.5555558),
+            (5.5555558, 55.5555573, 5.5555558),
+            (16.6666679, 66.6666641, 5.5555558),
+            (27.7777786, 77.7777786, 5.5555558),
+            (38.8888893, 88.8888855, 5.5555558),
+            (50.0000000, 100.0000000, 5.5555558),
+            (-50.0000000, 0.0000000, 16.6666679),
+            (-38.8888893, 11.1111107, 16.6666679),
+            (-27.7777786, 22.2222214, 16.6666679),
+            (-16.6666679, 33.3333321, 16.6666679),
+            (-5.5555558, 44.4444427, 16.6666679),
+            (5.5555558, 55.5555573, 16.6666679),
+            (16.6666679, 66.6666641, 16.6666679),
+            (27.7777786, 77.7777786, 16.6666679),
+            (38.8888893, 88.8888855, 16.6666679),
+            (50.0000000, 100.0000000, 16.6666679),
+            (-50.0000000, 0.0000000, 27.7777786),
+            (-38.8888893, 11.1111107, 27.7777786),
+            (-27.7777786, 22.2222214, 27.7777786),
+            (-16.6666679, 33.3333321, 27.7777786),
+            (-5.5555558, 44.4444427, 27.7777786),
+            (5.5555558, 55.5555573, 27.7777786),
+            (16.6666679, 66.6666641, 27.7777786),
+            (27.7777786, 77.7777786, 27.7777786),
+            (38.8888893, 88.8888855, 27.7777786),
+            (50.0000000, 100.0000000, 27.7777786),
+            (-50.0000000, 0.0000000, 38.8888893),
+            (-38.8888893, 11.1111107, 38.8888893),
+            (-27.7777786, 22.2222214, 38.8888893),
+            (-16.6666679, 33.3333321, 38.8888893),
+            (-5.5555558, 44.4444427, 38.8888893),
+            (5.5555558, 55.5555573, 38.8888893),
+            (16.6666679, 66.6666641, 38.8888893),
+            (27.7777786, 77.7777786, 38.8888893),
+            (38.8888893, 88.8888855, 38.8888893),
+            (50.0000000, 100.0000000, 38.8888893),
+            (-50.0000000, 0.0000000, 50.0000000),
+            (-38.8888893, 11.1111107, 50.0000000),
+            (-27.7777786, 22.2222214, 50.0000000),
+            (-16.6666679, 33.3333321, 50.0000000),
+            (-5.5555558, 44.4444427, 50.0000000),
+            (5.5555558, 55.5555573, 50.0000000),
+            (16.6666679, 66.6666641, 50.0000000),
+            (27.7777786, 77.7777786, 50.0000000),
+            (38.8888893, 88.8888855, 50.0000000),
+            (50.0000000, 100.0000000, 50.0000000),
+        ]
+
+        self.faces = [
+            1, 12, 2,
+            1, 11, 12,
+            2, 12, 3,
+            12, 13, 3,
+            3, 14, 4,
+            3, 13, 14,
+            4, 14, 5,
+            14, 15, 5,
+            5, 16, 6,
+            5, 15, 16,
+            6, 16, 7,
+            16, 17, 7,
+            7, 18, 8,
+            7, 17, 18,
+            8, 18, 9,
+            18, 19, 9,
+            9, 20, 10,
+            9, 19, 20,
+            11, 21, 12,
+            21, 22, 12,
+            12, 23, 13,
+            12, 22, 23,
+            13, 23, 14,
+            23, 24, 14,
+            14, 25, 15,
+            14, 24, 25,
+            15, 25, 16,
+            25, 26, 16,
+            16, 27, 17,
+            16, 26, 27,
+            17, 27, 18,
+            27, 28, 18,
+            18, 29, 19,
+            18, 28, 29,
+            19, 29, 20,
+            29, 30, 20,
+            21, 32, 22,
+            21, 31, 32,
+            22, 32, 23,
+            32, 33, 23,
+            23, 34, 24,
+            23, 33, 34,
+            24, 34, 25,
+            34, 35, 25,
+            25, 36, 26,
+            25, 35, 36,
+            26, 36, 27,
+            36, 37, 27,
+            27, 38, 28,
+            27, 37, 38,
+            28, 38, 29,
+            38, 39, 29,
+            29, 40, 30,
+            29, 39, 40,
+            31, 41, 32,
+            41, 42, 32,
+            32, 43, 33,
+            32, 42, 43,
+            33, 43, 34,
+            43, 44, 34,
+            34, 45, 35,
+            34, 44, 45,
+            35, 45, 36,
+            45, 46, 36,
+            36, 47, 37,
+            36, 46, 47,
+            37, 47, 38,
+            47, 48, 38,
+            38, 49, 39,
+            38, 48, 49,
+            39, 49, 40,
+            49, 50, 40,
+            41, 52, 42,
+            41, 51, 52,
+            42, 52, 43,
+            52, 53, 43,
+            43, 54, 44,
+            43, 53, 54,
+            44, 54, 45,
+            54, 55, 45,
+            45, 56, 46,
+            45, 55, 56,
+            46, 56, 47,
+            56, 57, 47,
+            47, 58, 48,
+            47, 57, 58,
+            48, 58, 49,
+            58, 59, 49,
+            49, 60, 50,
+            49, 59, 60,
+            51, 61, 52,
+            61, 62, 52,
+            52, 63, 53,
+            52, 62, 63,
+            53, 63, 54,
+            63, 64, 54,
+            54, 65, 55,
+            54, 64, 65,
+            55, 65, 56,
+            65, 66, 56,
+            56, 67, 57,
+            56, 66, 67,
+            57, 67, 58,
+            67, 68, 58,
+            58, 69, 59,
+            58, 68, 69,
+            59, 69, 60,
+            69, 70, 60,
+            61, 72, 62,
+            61, 71, 72,
+            62, 72, 63,
+            72, 73, 63,
+            63, 74, 64,
+            63, 73, 74,
+            64, 74, 65,
+            74, 75, 65,
+            65, 76, 66,
+            65, 75, 76,
+            66, 76, 67,
+            76, 77, 67,
+            67, 78, 68,
+            67, 77, 78,
+            68, 78, 69,
+            78, 79, 69,
+            69, 80, 70,
+            69, 79, 80,
+            71, 81, 72,
+            81, 82, 72,
+            72, 83, 73,
+            72, 82, 83,
+            73, 83, 74,
+            83, 84, 74,
+            74, 85, 75,
+            74, 84, 85,
+            75, 85, 76,
+            85, 86, 76,
+            76, 87, 77,
+            76, 86, 87,
+            77, 87, 78,
+            87, 88, 78,
+            78, 89, 79,
+            78, 88, 89,
+            79, 89, 80,
+            89, 90, 80,
+            81, 92, 82,
+            81, 91, 92,
+            82, 92, 83,
+            92, 93, 83,
+            83, 94, 84,
+            83, 93, 94,
+            84, 94, 85,
+            94, 95, 85,
+            85, 96, 86,
+            85, 95, 96,
+            86, 96, 87,
+            96, 97, 87,
+            87, 98, 88,
+            87, 97, 98,
+            88, 98, 89,
+            98, 99, 89,
+            89, 100, 90,
+            89, 99, 100
+        ]
+
+        self.coloring = [
+            [9, 12, 17, 24, 31, 38, 43, 46, 50, 62, 65, 68, 80, 84, 89, 92],
+            [6, 20, 25, 32, 37, 44, 51, 56, 59, 63, 70, 75, 82, 88, 90, 94, 96],
+            [2, 8, 10, 14, 26, 29, 33, 40, 48, 52, 55, 67, 73, 79, 86, 91, 98],
+            [4, 11, 16, 23, 28, 30, 35, 42, 49, 54, 57, 71, 74, 76, 78, 93, 97],
+            [3, 15, 18, 22, 34, 36, 39, 41, 53, 58, 60, 66, 72, 85, 99, 0, 87],
+            [7, 21, 27, 45, 47, 61, 64, 69, 77, 81, 83, 95, 1, 5, 13, 19],
+        ]
+        # fmt: on
+
+        self.dt = 1 / 60
+        self.num_test_frames = 100
+        self.num_substeps = 10
+        self.iterations = 10
+
+        stiffness = 1e5
+        kd = 1.0e-7
+
+        self.input_scale_factor = 1.0
+        self.renderer_scale_factor = 0.01
+        vertices = [wp.vec3(v) * self.input_scale_factor for v in self.pts]
+        fs_flatten = [fv - 1 for fv in self.faces]
+
+        builder = wp.sim.ModelBuilder()
+        builder.add_cloth_mesh(
+            pos=wp.vec3(0.0, 200.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=fs_flatten,
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            density=0.02,
+            tri_ke=stiffness,
+            tri_ka=stiffness,
+            tri_kd=kd,
+        )
+
+        self.model = builder.finalize(device=device)
+        self.model.ground = True
+        self.model.gravity = wp.vec3(0, -1000.0, 0)
+
+        self.model.soft_contact_ke = 1.0e4
+        self.model.soft_contact_kd = 1.0e2
+
+        coloring_wp = []
+        for color in self.coloring:
+            coloring_wp.append(wp.array(color, dtype=wp.int32, device=self.model.device))
+        self.model.coloring = coloring_wp
+
+        self.dt = self.dt / self.num_substeps
+        self.fixed_particles = [0, 9]
+
+        self.set_points_fixed(self.model, self.fixed_particles)
+
+        self.integrator = wp.sim.VBDIntegrator(self.model, self.iterations)
+        self.state0 = self.model.state()
+        self.state1 = self.model.state()
+
+        self.init_pos = np.array(self.state0.particle_q.numpy(), copy=True)
+
+    def run(self, test):
+        for _step in range(self.num_substeps * self.num_test_frames):
+            self.integrator.simulate(self.model, self.state0, self.state1, self.dt, None)
+            (self.state0, self.state1) = (self.state1, self.state0)
+
+        # examine that the simulation does not explode
+        final_pos = self.state0.particle_q.numpy()
+        test.assertTrue((final_pos < 1e5).all())
+        # examine that the simulation have moved
+        test.assertTrue((self.init_pos != final_pos).any())
+
+    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))
+
+            model.particle_flags = wp.array(flags, device=model.device)
+
+
+def test_vbd_cloth(test, device):
+    example = VBDClothSim(device)
+    example.run(test)
+
+
+devices = get_test_devices()
+
+
+class TestVBD(unittest.TestCase):
+    pass
+
+
+add_function_test(TestVBD, "test_vbd_cloth", test_vbd_cloth, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_vec.py

@@ -1035,6 +1035,142 @@ def test_casting_constructors(test, device, dtype, register_kernels=False):
     assert_np_equal(out, a_grad.numpy())
 
 
+def test_vec_assign(test, device, dtype, register_kernels=False):
+    np_type = np.dtype(dtype)
+    wp_type = wp.types.np_dtype_to_warp_type[np_type]
+
+    vec2 = wp.types.vector(length=2, dtype=wp_type)
+    vec3 = wp.types.vector(length=3, dtype=wp_type)
+    vec4 = wp.types.vector(length=4, dtype=wp_type)
+
+    def vectest_read_write_store(
+        x: wp.array(dtype=wp_type), a: wp.array(dtype=vec2), b: wp.array(dtype=vec3), c: wp.array(dtype=vec4)
+    ):
+        tid = wp.tid()
+
+        t = a[tid]
+        t[0] = x[tid]
+        a[tid] = t
+
+        u = b[tid]
+        u[1] = x[tid]
+        b[tid] = u
+
+        v = c[tid]
+        v[2] = x[tid]
+        c[tid] = v
+
+    def vectest_in_register(
+        x: wp.array(dtype=wp_type), y: wp.array(dtype=vec3), a: wp.array(dtype=vec2), b: wp.array(dtype=vec3)
+    ):
+        tid = wp.tid()
+
+        f = vec3(wp_type(0.0))
+        b_vec = b[tid]
+        f[0] = b_vec[1]
+        f[2] = b_vec[0] * b_vec[1]
+        y[tid] = f
+
+        g = wp_type(0.0)
+        a_vec = a[tid]
+        g = a_vec[0] + a_vec[1]
+        x[tid] = g
+
+    def vectest_in_register_overwrite(x: wp.array(dtype=vec3), a: wp.array(dtype=vec3)):
+        tid = wp.tid()
+
+        f = vec3(wp_type(0.0))
+        a_vec = a[tid]
+        f = a_vec
+        f[1] = wp_type(3.0)
+
+        x[tid] = f
+
+    def vectest_component(x: wp.array(dtype=vec3), y: wp.array(dtype=wp_type)):
+        i = wp.tid()
+
+        a = vec3(wp_type(0.0))
+        a.x = wp_type(1.0) * y[i]
+        a.y = wp_type(2.0) * y[i]
+        a.z = wp_type(3.0) * y[i]
+        x[i] = a
+
+    kernel_read_write_store = getkernel(vectest_read_write_store, suffix=dtype.__name__)
+    kernel_in_register = getkernel(vectest_in_register, suffix=dtype.__name__)
+    kernel_in_register_overwrite = getkernel(vectest_in_register_overwrite, suffix=dtype.__name__)
+    kernel_component = getkernel(vectest_component, suffix=dtype.__name__)
+
+    if register_kernels:
+        return
+
+    a = wp.ones(1, dtype=vec2, device=device, requires_grad=True)
+    b = wp.ones(1, dtype=vec3, device=device, requires_grad=True)
+    c = wp.ones(1, dtype=vec4, device=device, requires_grad=True)
+    x = wp.full(1, value=2.0, dtype=wp_type, device=device, requires_grad=True)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(kernel_read_write_store, dim=1, inputs=[x, a, b, c], device=device)
+
+    tape.backward(
+        grads={
+            a: wp.ones_like(a, requires_grad=False),
+            b: wp.ones_like(b, requires_grad=False),
+            c: wp.ones_like(c, requires_grad=False),
+        }
+    )
+
+    assert_np_equal(a.numpy(), np.array([[2.0, 1.0]], dtype=np_type))
+    assert_np_equal(b.numpy(), np.array([[1.0, 2.0, 1.0]], dtype=np_type))
+    assert_np_equal(c.numpy(), np.array([[1.0, 1.0, 2.0, 1.0]], dtype=np_type))
+    assert_np_equal(x.grad.numpy(), np.array([3.0], dtype=np_type))
+
+    tape.reset()
+
+    a = wp.ones(1, dtype=vec2, device=device, requires_grad=True)
+    b = wp.ones(1, dtype=vec3, device=device, requires_grad=True)
+    x = wp.zeros(1, dtype=wp_type, device=device, requires_grad=True)
+    y = wp.zeros(1, dtype=vec3, device=device, requires_grad=True)
+
+    with tape:
+        wp.launch(kernel_in_register, dim=1, inputs=[x, y, a, b], device=device)
+
+    tape.backward(grads={x: wp.ones_like(x, requires_grad=False), y: wp.ones_like(y, requires_grad=False)})
+
+    assert_np_equal(x.numpy(), np.array([2.0], dtype=np_type))
+    assert_np_equal(y.numpy(), np.array([[1.0, 0.0, 1.0]], dtype=np_type))
+    assert_np_equal(a.grad.numpy(), np.array([[1.0, 1.0]], dtype=np_type))
+    assert_np_equal(b.grad.numpy(), np.array([[1.0, 2.0, 0.0]], dtype=np_type))
+
+    tape.reset()
+
+    x = wp.zeros(1, dtype=vec3, device=device, requires_grad=True)
+    y = wp.ones(1, dtype=wp_type, device=device, requires_grad=True)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(kernel_component, dim=1, inputs=[x, y], device=device)
+
+    tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
+
+    assert_np_equal(x.numpy(), np.array([[1.0, 2.0, 3.0]], dtype=np_type))
+    assert_np_equal(y.grad.numpy(), np.array([6.0], dtype=np_type))
+
+    tape.reset()
+
+    x = wp.zeros(1, dtype=vec3, device=device, requires_grad=True)
+    a = wp.ones(1, dtype=vec3, device=device, requires_grad=True)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(kernel_in_register_overwrite, dim=1, inputs=[x, a], device=device)
+
+    tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
+
+    assert_np_equal(x.numpy(), np.array([[1.0, 3.0, 1.0]], dtype=np_type))
+    assert_np_equal(a.grad.numpy(), np.array([[1.0, 0.0, 1.0]], dtype=np_type))
+
+
 @wp.kernel
 def test_vector_constructor_value_func():
     a = wp.vec2()
@@ -1170,6 +1306,13 @@ for dtype in np_float_types:
         devices=devices,
         dtype=dtype,
     )
+    add_function_test_register_kernel(
+        TestVec,
+        f"test_vec_assign_{dtype.__name__}",
+        test_vec_assign,
+        devices=devices,
+        dtype=dtype,
+    )
 
 add_function_test(
     TestVec,