warp-lang 1.7.2rc1__py3-none-macosx_10_13_universal2.whl → 1.8.1__py3-none-macosx_10_13_universal2.whl

warp/tests/test_spatial.py

@@ -2030,6 +2030,486 @@ def test_transform_to_matrix(test, device, dtype, register_kernels=False):
     wp.launch(kernel, dim=1, device=device)
 
 
+def test_transform_compose(test, device, dtype, register_kernels=False):
+    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    mat44 = wp.types.matrix((4, 4), wptype)
+    vec3 = wp.types.vector(3, wptype)
+    quat = wp.types.quaternion(wptype)
+
+    def transform_compose_kernel():
+        p = vec3(wptype(1.0), wptype(2.0), wptype(3.0))
+        q = quat(wptype(-0.4), wptype(0.2), wptype(-0.4), wptype(0.8))
+        s = vec3(wptype(4.0), wptype(5.0), wptype(6.0))
+        m = wp.transform_compose(p, q, s)
+        # fmt: off
+        wp.expect_near(
+            m,
+            mat44(
+                wptype(0.6 * 4.0), wptype(0.48 * 5.0), wptype(0.64 * 6.0), wptype(1.0),
+                wptype(-0.8 * 4.0), wptype(0.36 * 5.0), wptype(0.48 * 6.0), wptype(2.0),
+                wptype(0.0 * 4.0), wptype(-0.8 * 5.0), wptype(0.6 * 6.0), wptype(3.0),
+                wptype(0.0), wptype(0.0), wptype(0.0), wptype(1.0),
+            ),
+            tolerance=wptype(1e-2),
+        )
+        # fmt: on
+
+    kernel = getkernel(transform_compose_kernel, suffix=dtype.__name__)
+
+    if register_kernels:
+        return
+
+    wp.launch(kernel, dim=1, device=device)
+
+
+def test_transform_decompose(test, device, dtype, register_kernels=False):
+    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    mat44 = wp.types.matrix((4, 4), wptype)
+    vec3 = wp.types.vector(3, wptype)
+    quat = wp.types.quaternion(wptype)
+
+    def transform_decompose_kernel():
+        # fmt: off
+        m = mat44(
+            wptype(0.6 * 4.0), wptype(0.48 * 5.0), wptype(0.64 * 6.0), wptype(1.0),
+            wptype(-0.8 * 4.0), wptype(0.36 * 5.0), wptype(0.48 * 6.0), wptype(2.0),
+            wptype(0.0 * 4.0), wptype(-0.8 * 5.0), wptype(0.6 * 6.0), wptype(3.0),
+            wptype(0.0), wptype(0.0), wptype(0.0), wptype(1.0),
+        )
+        # fmt: on
+        p, q, s = wp.transform_decompose(m)
+        wp.expect_near(p, vec3(wptype(1.0), wptype(2.0), wptype(3.0)), tolerance=wptype(1e-2))
+        wp.expect_near(q, quat(wptype(-0.4), wptype(0.2), wptype(-0.4), wptype(0.8)), tolerance=wptype(1e-2))
+        wp.expect_near(s, vec3(wptype(4.0), wptype(5.0), wptype(6.0)), tolerance=wptype(1e-2))
+
+    kernel = getkernel(transform_decompose_kernel, suffix=dtype.__name__)
+
+    if register_kernels:
+        return
+
+    wp.launch(kernel, dim=1, device=device)
+
+
+def test_transform_getter_setter(test, device):
+    t = wp.transform()
+
+    a = wp.vec3(1.0, 2.0, 3.0)
+    b = wp.quat(0.0, 0.0, 0.0, 1.0)
+
+    t.p = a
+    t.q = b
+
+    c = t.p
+    d = t.q
+
+    test.assertEqual(c, a)
+    test.assertEqual(d, b)
+
+
+@wp.kernel
+def transform_extract_subscript(x: wp.array(dtype=wp.transform), y: wp.array(dtype=float)):
+    tid = wp.tid()
+
+    a = x[tid]
+    b = a[0] + 2.0 * a[1] + 3.0 * a[2] + 4.0 * a[3] + 5.0 * a[4] + 6.0 * a[5] + 7.0 * a[6]
+    y[tid] = b
+
+
+@wp.kernel
+def transform_extract_attribute(
+    x: wp.array(dtype=wp.transform), y: wp.array(dtype=wp.vec3), z: wp.array(dtype=wp.quat)
+):
+    tid = wp.tid()
+
+    a = x[tid]
+    p = a.p
+    q = a.q
+
+    y[tid] = p
+    z[tid] = q
+
+
+def test_transform_extract(test, device):
+    x = wp.ones(1, dtype=wp.transform, requires_grad=True, device=device)
+    y = wp.zeros(1, dtype=float, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(transform_extract_subscript, 1, inputs=[x], outputs=[y], device=device)
+
+    y.grad = wp.ones_like(y)
+
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.array([28.0], dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]], dtype=float))
+
+    x = wp.ones(1, dtype=wp.transform, requires_grad=True, device=device)
+    y = wp.zeros(1, dtype=wp.vec3, requires_grad=True, device=device)
+    z = wp.zeros(1, dtype=wp.quat, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(transform_extract_attribute, 1, inputs=[x], outputs=[y, z], device=device)
+
+    y.grad = wp.ones_like(y)
+    z.grad = wp.ones_like(z)
+
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.array([[1.0, 1.0, 1.0]], dtype=float))
+    assert_np_equal(z.numpy(), np.array([[1.0, 1.0, 1.0, 1.0]], dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array([[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]], dtype=float))
+
+
+@wp.kernel
+def transform_assign_subscript(x: wp.array(dtype=float), y: wp.array(dtype=wp.transform)):
+    i = wp.tid()
+
+    a = wp.transform()
+    a[0] = 1.0 * x[i]
+    a[1] = 2.0 * x[i]
+    a[2] = 3.0 * x[i]
+    a[3] = 4.0 * x[i]
+    a[4] = 5.0 * x[i]
+    a[5] = 6.0 * x[i]
+    a[6] = 7.0 * x[i]
+    y[i] = a
+
+
+@wp.kernel
+def transform_assign_attribute(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=wp.quat), z: wp.array(dtype=wp.transform)):
+    i = wp.tid()
+
+    a = wp.transform()
+    a.p = x[i]
+    a.q = y[i]
+    z[i] = a
+
+
+def test_transform_assign(test, device):
+    x = wp.ones(1, dtype=float, requires_grad=True, device=device)
+    y = wp.zeros(1, dtype=wp.transform, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(transform_assign_subscript, 1, inputs=[x], outputs=[y], device=device)
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]], dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array([28.0], dtype=float))
+
+    x = wp.ones(1, dtype=wp.vec3, requires_grad=True, device=device)
+    y = wp.ones(1, dtype=wp.quat, requires_grad=True, device=device)
+    z = wp.zeros(1, dtype=wp.transform, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(transform_assign_attribute, 1, inputs=[x, y], outputs=[z], device=device)
+
+    z.grad = wp.ones_like(z)
+    tape.backward()
+
+    assert_np_equal(z.numpy(), np.array([[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]], dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array([[1.0, 1.0, 1.0]], dtype=float))
+    assert_np_equal(y.grad.numpy(), np.array([[1.0, 1.0, 1.0, 1.0]], dtype=float))
+
+
+def test_transform_assign_copy(test, device):
+    saved_enable_vector_component_overwrites_setting = wp.config.enable_vector_component_overwrites
+    try:
+        wp.config.enable_vector_component_overwrites = True
+
+        @wp.kernel
+        def transform_assign_overwrite(x: wp.array(dtype=wp.transform), y: wp.array(dtype=wp.transform)):
+            tid = wp.tid()
+
+            a = wp.transform()
+            b = x[tid]
+            a = b
+            a[1] = 3.0
+
+            y[tid] = a
+
+        x = wp.ones(1, dtype=wp.transform, device=device, requires_grad=True)
+        y = wp.zeros(1, dtype=wp.transform, device=device, requires_grad=True)
+
+        tape = wp.Tape()
+        with tape:
+            wp.launch(transform_assign_overwrite, dim=1, inputs=[x, y], device=device)
+
+        y.grad = wp.ones_like(y, requires_grad=False)
+        tape.backward()
+
+        assert_np_equal(y.numpy(), np.array([[1.0, 3.0, 1.0, 1.0, 1.0, 1.0, 1.0]], dtype=float))
+        assert_np_equal(x.grad.numpy(), np.array([[1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0]], dtype=float))
+
+    finally:
+        wp.config.enable_vector_component_overwrites = saved_enable_vector_component_overwrites_setting
+
+
+@wp.kernel
+def transform_array_extract_subscript(x: wp.array2d(dtype=wp.transform), y: wp.array2d(dtype=float)):
+    i, j = wp.tid()
+    a = x[i, j][0]
+    b = x[i, j][1]
+    c = x[i, j][2]
+    d = x[i, j][3]
+    e = x[i, j][4]
+    f = x[i, j][5]
+    g = x[i, j][6]
+    y[i, j] = 1.0 * a + 2.0 * b + 3.0 * c + 4.0 * d + 5.0 * e + 6.0 * f + 7.0 * g
+
+
+@wp.kernel
+def transform_array_extract_attribute(
+    x: wp.array2d(dtype=wp.transform), y: wp.array2d(dtype=wp.vec3), z: wp.array2d(dtype=wp.quat)
+):
+    i, j = wp.tid()
+    a = x[i, j].p
+    b = x[i, j].q
+    y[i, j] = a
+    z[i, j] = b
+
+
+def test_transform_array_extract(test, device):
+    x = wp.ones((1, 1), dtype=wp.transform, requires_grad=True, device=device)
+    y = wp.zeros((1, 1), dtype=float, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(transform_array_extract_subscript, (1, 1), inputs=[x], outputs=[y], device=device)
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.array([[28.0]], dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array([[[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]]], dtype=float))
+
+    x = wp.ones((1, 1), dtype=wp.transform, requires_grad=True, device=device)
+    y = wp.zeros((1, 1), dtype=wp.vec3, requires_grad=True, device=device)
+    z = wp.zeros((1, 1), dtype=wp.quat, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(transform_array_extract_attribute, (1, 1), inputs=[x], outputs=[y, z], device=device)
+
+    y.grad = wp.ones_like(y)
+    z.grad = wp.ones_like(z)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.array([[[1.0, 1.0, 1.0]]], dtype=float))
+    assert_np_equal(z.numpy(), np.array([[[1.0, 1.0, 1.0, 1.0]]], dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array([[[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]]], dtype=float))
+
+
+@wp.kernel
+def transform_array_assign_subscript(x: wp.array2d(dtype=float), y: wp.array2d(dtype=wp.transform)):
+    i, j = wp.tid()
+
+    y[i, j][0] = 1.0 * x[i, j]
+    y[i, j][1] = 2.0 * x[i, j]
+    y[i, j][2] = 3.0 * x[i, j]
+    y[i, j][3] = 4.0 * x[i, j]
+    y[i, j][4] = 5.0 * x[i, j]
+    y[i, j][5] = 6.0 * x[i, j]
+    y[i, j][6] = 7.0 * x[i, j]
+
+
+# @wp.kernel
+# def transform_array_assign_attribute(x: wp.array2d(dtype=wp.vec3), y: wp.array2d(dtype=wp.quat), z: wp.array2d(dtype=wp.transform)):
+#     i, j = wp.tid()
+
+#     z[i, j].p = x[i, j]
+#     z[i, j].q = y[i, j]
+
+
+def test_transform_array_assign(test, device):
+    x = wp.ones((1, 1), dtype=float, requires_grad=True, device=device)
+    y = wp.zeros((1, 1), dtype=wp.transform, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(transform_array_assign_subscript, (1, 1), inputs=[x], outputs=[y], device=device)
+
+    assert_np_equal(y.numpy(), np.array([[[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]]], dtype=float))
+    # TODO: gradient propagation for in-place array assignment
+
+    x = wp.ones((1, 1), dtype=wp.vec3, requires_grad=True, device=device)
+    y = wp.ones((1, 1), dtype=wp.quat, requires_grad=True, device=device)
+    z = wp.zeros((1, 1), dtype=wp.transform, requires_grad=True, device=device)
+
+    # TODO: transform_array_assign_attribute
+    # tape = wp.Tape()
+    # with tape:
+    #     wp.launch(transform_array_assign_attribute, (1, 1), inputs=[x, y], outputs=[z], device=device)
+
+    # assert_np_equal(z.numpy(), np.array([[[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]]], dtype=float))
+    # TODO: gradient propagation for in-place array assignment
+
+
+@wp.kernel
+def transform_add_inplace_subscript(x: wp.array(dtype=wp.transform), y: wp.array(dtype=wp.transform)):
+    i = wp.tid()
+
+    a = wp.transform()
+    b = x[i]
+
+    a[0] += 1.0 * b[0]
+    a[1] += 2.0 * b[1]
+    a[2] += 3.0 * b[2]
+    a[3] += 4.0 * b[3]
+    a[4] += 5.0 * b[4]
+    a[5] += 6.0 * b[5]
+    a[6] += 7.0 * b[6]
+
+    y[i] = a
+
+
+@wp.kernel
+def transform_add_inplace_attribute(x: wp.array(dtype=wp.transform), y: wp.array(dtype=wp.transform)):
+    i = wp.tid()
+
+    a = wp.transform()
+    b = x[i]
+
+    a.p += b.p
+
+    y[i] = a
+
+
+def test_transform_add_inplace(test, device):
+    x = wp.ones(1, dtype=wp.transform, requires_grad=True, device=device)
+    y = wp.zeros(1, dtype=wp.transform, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(transform_add_inplace_subscript, 1, inputs=[x], outputs=[y], device=device)
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]], dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]], dtype=float))
+
+    x = wp.ones(1, dtype=wp.transform, requires_grad=True, device=device)
+    y = wp.zeros(1, dtype=wp.transform, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(transform_add_inplace_attribute, 1, inputs=[x], outputs=[y], device=device)
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.array([[1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0]], dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array([[1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0]], dtype=float))
+
+
+@wp.kernel
+def transform_sub_inplace_subscript(x: wp.array(dtype=wp.transform), y: wp.array(dtype=wp.transform)):
+    i = wp.tid()
+
+    a = wp.transform()
+    b = x[i]
+
+    a[0] -= 1.0 * b[0]
+    a[1] -= 2.0 * b[1]
+    a[2] -= 3.0 * b[2]
+    a[3] -= 4.0 * b[3]
+    a[4] -= 5.0 * b[4]
+    a[5] -= 6.0 * b[5]
+    a[6] -= 7.0 * b[6]
+
+    y[i] = a
+
+
+@wp.kernel
+def transform_sub_inplace_attribute(x: wp.array(dtype=wp.transform), y: wp.array(dtype=wp.transform)):
+    i = wp.tid()
+
+    a = wp.transform()
+    b = x[i]
+
+    a.p -= b.p
+
+    y[i] = a
+
+
+def test_transform_sub_inplace(test, device):
+    x = wp.ones(1, dtype=wp.transform, requires_grad=True, device=device)
+    y = wp.zeros(1, dtype=wp.transform, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(transform_sub_inplace_subscript, 1, inputs=[x], outputs=[y], device=device)
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.array([[-1.0, -2.0, -3.0, -4.0, -5.0, -6.0, -7.0]], dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array([[-1.0, -2.0, -3.0, -4.0, -5.0, -6.0, -7.0]], dtype=float))
+
+    x = wp.ones(1, dtype=wp.transform, requires_grad=True, device=device)
+    y = wp.zeros(1, dtype=wp.transform, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(transform_sub_inplace_attribute, 1, inputs=[x], outputs=[y], device=device)
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.array([[-1.0, -1.0, -1.0, 0.0, 0.0, 0.0, 0.0]], dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array([[-1.0, -1.0, -1.0, 0.0, 0.0, 0.0, 0.0]], dtype=float))
+
+
+@wp.kernel
+def transform_array_add_inplace(x: wp.array(dtype=wp.transform), y: wp.array(dtype=wp.transform)):
+    i = wp.tid()
+
+    y[i] += x[i]
+
+
+def test_transform_array_add_inplace(test, device):
+    x = wp.ones(1, dtype=wp.transform, requires_grad=True, device=device)
+    y = wp.zeros(1, dtype=wp.transform, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(transform_array_add_inplace, 1, inputs=[x], outputs=[y], device=device)
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.array([[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]], dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array([[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]], dtype=float))
+
+
+@wp.kernel
+def transform_array_sub_inplace(x: wp.array(dtype=wp.transform), y: wp.array(dtype=wp.transform)):
+    i = wp.tid()
+
+    y[i] -= x[i]
+
+
+def test_transform_array_sub_inplace(test, device):
+    x = wp.ones(1, dtype=wp.transform, requires_grad=True, device=device)
+    y = wp.zeros(1, dtype=wp.transform, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(transform_array_sub_inplace, 1, inputs=[x], outputs=[y], device=device)
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.array([[-1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0]], dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array([[-1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0]], dtype=float))
+
+
 devices = get_test_devices()
 
 
@@ -2220,9 +2700,35 @@ for dtype in np_float_types:
         devices=devices,
         dtype=dtype,
     )
-
+    add_function_test_register_kernel(
+        TestSpatial,
+        f"test_transform_compose_{dtype.__name__}",
+        test_transform_compose,
+        devices=devices,
+        dtype=dtype,
+    )
+    add_function_test_register_kernel(
+        TestSpatial,
+        f"test_transform_decompose_{dtype.__name__}",
+        test_transform_decompose,
+        devices=devices,
+        dtype=dtype,
+    )
     # \TODO: test spatial_mass and spatial_jacobian
 
+add_function_test(
+    TestSpatial, "test_transform_getter_setter", test_transform_getter_setter, devices=wp.get_device("cpu")
+)
+add_function_test(TestSpatial, "test_transform_extract", test_transform_extract, devices=devices)
+add_function_test(TestSpatial, "test_transform_assign", test_transform_assign, devices=devices)
+add_function_test(TestSpatial, "test_transform_assign_copy", test_transform_assign_copy, devices=devices)
+add_function_test(TestSpatial, "test_transform_array_extract", test_transform_array_extract, devices=devices)
+add_function_test(TestSpatial, "test_transform_array_assign", test_transform_array_assign, devices=devices)
+add_function_test(TestSpatial, "test_transform_add_inplace", test_transform_add_inplace, devices=devices)
+add_function_test(TestSpatial, "test_transform_sub_inplace", test_transform_sub_inplace, devices=devices)
+add_function_test(TestSpatial, "test_transform_array_add_inplace", test_transform_array_add_inplace, devices=devices)
+add_function_test(TestSpatial, "test_transform_array_sub_inplace", test_transform_array_sub_inplace, devices=devices)
+
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/test_static.py

@@ -559,6 +559,54 @@ def test_static_len_query(test, _):
     wp.launch(static_len_query_kernel, 1, inputs=(v1,))
 
 
+@wp.func
+def func_1() -> int:
+    return 1
+
+
+@wp.func
+def func_2() -> int:
+    return 2
+
+
+funcs = [func_1, func_2]
+
+
+def unresolved_builder(funcids):
+    _funcs = [funcs[id] for id in funcids]
+
+    @wp.kernel
+    def eval(input: wp.array(dtype=int), output: wp.array(dtype=int)):
+        for i in range(wp.static(len(_funcs))):
+            output[0] = wp.static(_funcs[i])()
+
+    return eval
+
+
+def test_unresolved_static_expression(test, device):
+    # The module hash will need to be updated from the static expressions
+    # resolved at code generation time, since some of them cannot be evaluated
+    # at declaration time.
+    with wp.ScopedDevice(device):
+        output1 = wp.array((1,), dtype=int)
+        wp.launch(
+            unresolved_builder([0]),
+            dim=(1,),
+            inputs=[wp.array(np.array([0]), dtype=int)],
+            outputs=[output1],
+        )
+        test.assertEqual(output1.numpy()[0], 1)
+
+        output2 = wp.array((1,), dtype=int)
+        wp.launch(
+            unresolved_builder([1]),
+            dim=(1,),
+            inputs=[wp.array(np.array([1]), dtype=int)],
+            outputs=[output2],
+        )
+        test.assertEqual(output2.numpy()[0], 2)
+
+
 devices = get_test_devices()
 
 

warp/tests/test_struct.py

@@ -768,8 +768,8 @@ class TestStruct(unittest.TestCase):
         v.value[2] = 3.0
 
         arr = wp.array([v], dtype=VecStruct)
-        expected = np.array(([1.0, 2.0, 3.0],))
-        assert np.all(arr.numpy().tolist() == expected)
+        expected = np.array([[[1.0, 2.0, 3.0]]])
+        np.testing.assert_equal(arr.numpy().tolist(), expected)
 
 
 add_function_test(TestStruct, "test_step", test_step, devices=devices)

warp/tests/test_tape.py

@@ -157,6 +157,43 @@ def test_tape_zero_multiple_outputs(test, device):
     assert_np_equal(x.grad.numpy(), np.ones(3, dtype=float))
 
 
+@wp.struct
+class NestedStruct:
+    arr: wp.array(dtype=float)
+
+
+@wp.struct
+class WrapperStruct:
+    nested: NestedStruct
+
+
+@wp.kernel
+def nested_loss_kernel(wrapper: WrapperStruct, loss: wp.array(dtype=float)):
+    i = wp.tid()
+    wp.atomic_add(loss, 0, wrapper.nested.arr[i])
+
+
+def test_tape_nested_struct(test, device):
+    wrapper = WrapperStruct()
+    wrapper.nested = NestedStruct()
+    wrapper.nested.arr = wp.ones(shape=(1,), dtype=float, requires_grad=True, device=device)
+
+    loss = wp.zeros(shape=(1,), dtype=float, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(nested_loss_kernel, dim=1, inputs=(wrapper, loss), device=device)
+
+    assert_np_equal(loss.numpy(), [1.0])
+
+    tape.backward(loss)
+    assert_np_equal(wrapper.nested.arr.grad.numpy(), [1.0])
+
+    tape.zero()
+
+    assert_np_equal(wrapper.nested.arr.grad.numpy(), [0.0])
+
+
 def test_tape_visualize(test, device):
     dim = 8
     tape = wp.Tape()
@@ -196,6 +233,7 @@ add_function_test(TestTape, "test_tape_mul_constant", test_tape_mul_constant, de
 add_function_test(TestTape, "test_tape_mul_variable", test_tape_mul_variable, devices=devices)
 add_function_test(TestTape, "test_tape_dot_product", test_tape_dot_product, devices=devices)
 add_function_test(TestTape, "test_tape_zero_multiple_outputs", test_tape_zero_multiple_outputs, devices=devices)
+add_function_test(TestTape, "test_tape_nested_struct", test_tape_nested_struct, devices=devices)
 add_function_test(TestTape, "test_tape_visualize", test_tape_visualize, devices=devices)