warp-lang 1.8.0__py3-none-manylinux_2_34_aarch64.whl → 1.8.1__py3-none-manylinux_2_34_aarch64.whl

warp/tests/test_atomic_cas.py

@@ -19,54 +19,63 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
+kernel_cache = {}
+
+
+def getkernel(func, suffix=""):
+    key = func.__name__ + "_" + suffix
+    if key not in kernel_cache:
+        kernel_cache[key] = wp.Kernel(func=func, key=key)
+    return kernel_cache[key]
+
+
+def test_atomic_cas(test, device, dtype, register_kernels=False):
+    warp_type = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    n = 100
+    counter = wp.array([0], dtype=warp_type, device=device)
+    lock = wp.array([0], dtype=warp_type, device=device)
 
-def create_spinlock_test(dtype):
     @wp.func
-    def spinlock_acquire(lock: wp.array(dtype=dtype)):
+    def spinlock_acquire_1d(lock: wp.array(dtype=warp_type)):
         # Try to acquire the lock by setting it to 1 if it's 0
-        while wp.atomic_cas(lock, 0, dtype(0), dtype(1)) == 1:
+        while wp.atomic_cas(lock, 0, warp_type(0), warp_type(1)) == 1:
             pass
 
     @wp.func
-    def spinlock_release(lock: wp.array(dtype=dtype)):
+    def spinlock_release_1d(lock: wp.array(dtype=warp_type)):
         # Release the lock by setting it back to 0
-        wp.atomic_exch(lock, 0, dtype(0))
+        wp.atomic_exch(lock, 0, warp_type(0))
 
     @wp.func
-    def volatile_read(ptr: wp.array(dtype=dtype), index: int):
-        value = wp.atomic_exch(ptr, index, dtype(0))
+    def volatile_read_1d(ptr: wp.array(dtype=warp_type), index: int):
+        value = wp.atomic_exch(ptr, index, warp_type(0))
         wp.atomic_exch(ptr, index, value)
         return value
 
-    @wp.kernel
-    def test_spinlock_counter(counter: wp.array(dtype=dtype), lock: wp.array(dtype=dtype)):
+    def test_spinlock_counter_1d(counter: wp.array(dtype=warp_type), lock: wp.array(dtype=warp_type)):
         # Try to acquire the lock
-        spinlock_acquire(lock)
+        spinlock_acquire_1d(lock)
 
         # Critical section - increment counter
         # counter[0] = counter[0] + 1 # This gives wrong results - counter should be marked as volatile
 
         # Work around since warp arrays cannot be marked as volatile
-        value = volatile_read(counter, 0)
-        counter[0] = value + dtype(1)
+        value = volatile_read_1d(counter, 0)
+        counter[0] = value + warp_type(1)
 
         # Release the lock
-        spinlock_release(lock)
-
-    return test_spinlock_counter
+        spinlock_release_1d(lock)
 
+    kernel = getkernel(test_spinlock_counter_1d, suffix=dtype.__name__)
 
-def test_atomic_cas(test, device, warp_type, numpy_type):
-    n = 100
-    counter = wp.array([0], dtype=warp_type, device=device)
-    lock = wp.array([0], dtype=warp_type, device=device)
+    if register_kernels:
+        return
 
-    test_spinlock_counter = create_spinlock_test(warp_type)
-    wp.launch(test_spinlock_counter, dim=n, inputs=[counter, lock], device=device)
+    wp.launch(kernel, dim=n, inputs=[counter, lock], device=device)
 
     # Verify counter reached n
     counter_np = counter.numpy()
-    expected = np.array([n], dtype=numpy_type)
+    expected = np.array([n], dtype=dtype)
 
     if not np.array_equal(counter_np, expected):
         print(f"Counter mismatch: expected {expected}, got {counter_np}")
@@ -74,53 +83,53 @@ def test_atomic_cas(test, device, warp_type, numpy_type):
     assert_np_equal(counter_np, expected)
 
 
-def create_spinlock_test_2d(dtype):
+def test_atomic_cas_2d(test, device, dtype, register_kernels=False):
+    warp_type = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    n = 100
+    counter = wp.array([0], dtype=warp_type, device=device)
+    lock = wp.zeros(shape=(1, 1), dtype=warp_type, device=device)
+
     @wp.func
-    def spinlock_acquire(lock: wp.array(dtype=dtype, ndim=2)):
+    def spinlock_acquire_2d(lock: wp.array2d(dtype=warp_type)):
         # Try to acquire the lock by setting it to 1 if it's 0
-        while wp.atomic_cas(lock, 0, 0, dtype(0), dtype(1)) == 1:
+        while wp.atomic_cas(lock, 0, 0, warp_type(0), warp_type(1)) == 1:
             pass
 
     @wp.func
-    def spinlock_release(lock: wp.array(dtype=dtype, ndim=2)):
+    def spinlock_release_2d(lock: wp.array2d(dtype=warp_type)):
         # Release the lock by setting it back to 0
-        wp.atomic_exch(lock, 0, 0, dtype(0))
+        wp.atomic_exch(lock, 0, 0, warp_type(0))
 
     @wp.func
-    def volatile_read(ptr: wp.array(dtype=dtype), index: int):
-        value = wp.atomic_exch(ptr, index, dtype(0))
+    def volatile_read_2d(ptr: wp.array(dtype=warp_type), index: int):
+        value = wp.atomic_exch(ptr, index, warp_type(0))
         wp.atomic_exch(ptr, index, value)
         return value
 
-    @wp.kernel
-    def test_spinlock_counter(counter: wp.array(dtype=dtype), lock: wp.array(dtype=dtype, ndim=2)):
+    def test_spinlock_counter_2d(counter: wp.array(dtype=warp_type), lock: wp.array2d(dtype=warp_type)):
         # Try to acquire the lock
-        spinlock_acquire(lock)
+        spinlock_acquire_2d(lock)
 
         # Critical section - increment counter
         # counter[0] = counter[0] + 1 # This gives wrong results - counter should be marked as volatile
 
         # Work around since warp arrays cannot be marked as volatile
-        value = volatile_read(counter, 0)
-        counter[0] = value + dtype(1)
+        value = volatile_read_2d(counter, 0)
+        counter[0] = value + warp_type(1)
 
         # Release the lock
-        spinlock_release(lock)
-
-    return test_spinlock_counter
+        spinlock_release_2d(lock)
 
+    kernel = getkernel(test_spinlock_counter_2d, suffix=dtype.__name__)
 
-def test_atomic_cas_2d(test, device, warp_type, numpy_type):
-    n = 100
-    counter = wp.array([0], dtype=warp_type, device=device)
-    lock = wp.zeros(shape=(1, 1), dtype=warp_type, device=device)
+    if register_kernels:
+        return
 
-    test_spinlock_counter = create_spinlock_test_2d(warp_type)
-    wp.launch(test_spinlock_counter, dim=n, inputs=[counter, lock], device=device)
+    wp.launch(kernel, dim=n, inputs=[counter, lock], device=device)
 
     # Verify counter reached n
     counter_np = counter.numpy()
-    expected = np.array([n], dtype=numpy_type)
+    expected = np.array([n], dtype=dtype)
 
     if not np.array_equal(counter_np, expected):
         print(f"Counter mismatch: expected {expected}, got {counter_np}")
@@ -128,53 +137,53 @@ def test_atomic_cas_2d(test, device, warp_type, numpy_type):
     assert_np_equal(counter_np, expected)
 
 
-def create_spinlock_test_3d(dtype):
+def test_atomic_cas_3d(test, device, dtype, register_kernels=False):
+    warp_type = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    n = 100
+    counter = wp.array([0], dtype=warp_type, device=device)
+    lock = wp.zeros(shape=(1, 1, 1), dtype=warp_type, device=device)
+
     @wp.func
-    def spinlock_acquire(lock: wp.array(dtype=dtype, ndim=3)):
+    def spinlock_acquire_3d(lock: wp.array3d(dtype=warp_type)):
         # Try to acquire the lock by setting it to 1 if it's 0
-        while wp.atomic_cas(lock, 0, 0, 0, dtype(0), dtype(1)) == 1:
+        while wp.atomic_cas(lock, 0, 0, 0, warp_type(0), warp_type(1)) == 1:
             pass
 
     @wp.func
-    def spinlock_release(lock: wp.array(dtype=dtype, ndim=3)):
+    def spinlock_release_3d(lock: wp.array3d(dtype=warp_type)):
         # Release the lock by setting it back to 0
-        wp.atomic_exch(lock, 0, 0, 0, dtype(0))
+        wp.atomic_exch(lock, 0, 0, 0, warp_type(0))
 
     @wp.func
-    def volatile_read(ptr: wp.array(dtype=dtype), index: int):
-        value = wp.atomic_exch(ptr, index, dtype(0))
+    def volatile_read_3d(ptr: wp.array(dtype=warp_type), index: int):
+        value = wp.atomic_exch(ptr, index, warp_type(0))
         wp.atomic_exch(ptr, index, value)
         return value
 
-    @wp.kernel
-    def test_spinlock_counter(counter: wp.array(dtype=dtype), lock: wp.array(dtype=dtype, ndim=3)):
+    def test_spinlock_counter_3d(counter: wp.array(dtype=warp_type), lock: wp.array3d(dtype=warp_type)):
         # Try to acquire the lock
-        spinlock_acquire(lock)
+        spinlock_acquire_3d(lock)
 
         # Critical section - increment counter
         # counter[0] = counter[0] + 1 # This gives wrong results - counter should be marked as volatile
 
         # Work around since warp arrays cannot be marked as volatile
-        value = volatile_read(counter, 0)
-        counter[0] = value + dtype(1)
+        value = volatile_read_3d(counter, 0)
+        counter[0] = value + warp_type(1)
 
         # Release the lock
-        spinlock_release(lock)
+        spinlock_release_3d(lock)
 
-    return test_spinlock_counter
+    kernel = getkernel(test_spinlock_counter_3d, suffix=dtype.__name__)
 
+    if register_kernels:
+        return
 
-def test_atomic_cas_3d(test, device, warp_type, numpy_type):
-    n = 100
-    counter = wp.array([0], dtype=warp_type, device=device)
-    lock = wp.zeros(shape=(1, 1, 1), dtype=warp_type, device=device)
-
-    test_spinlock_counter = create_spinlock_test_3d(warp_type)
-    wp.launch(test_spinlock_counter, dim=n, inputs=[counter, lock], device=device)
+    wp.launch(kernel, dim=n, inputs=[counter, lock], device=device)
 
     # Verify counter reached n
     counter_np = counter.numpy()
-    expected = np.array([n], dtype=numpy_type)
+    expected = np.array([n], dtype=dtype)
 
     if not np.array_equal(counter_np, expected):
         print(f"Counter mismatch: expected {expected}, got {counter_np}")
@@ -218,17 +227,53 @@ def create_spinlock_test_4d(dtype):
     return test_spinlock_counter
 
 
-def test_atomic_cas_4d(test, device, warp_type, numpy_type):
+def test_atomic_cas_4d(test, device, dtype, register_kernels=False):
+    warp_type = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
     n = 100
     counter = wp.array([0], dtype=warp_type, device=device)
     lock = wp.zeros(shape=(1, 1, 1, 1), dtype=warp_type, device=device)
 
-    test_spinlock_counter = create_spinlock_test_4d(warp_type)
-    wp.launch(test_spinlock_counter, dim=n, inputs=[counter, lock], device=device)
+    @wp.func
+    def spinlock_acquire_4d(lock: wp.array4d(dtype=warp_type)):
+        # Try to acquire the lock by setting it to 1 if it's 0
+        while wp.atomic_cas(lock, 0, 0, 0, 0, warp_type(0), warp_type(1)) == 1:
+            pass
+
+    @wp.func
+    def spinlock_release_4d(lock: wp.array4d(dtype=warp_type)):
+        # Release the lock by setting it back to 0
+        wp.atomic_exch(lock, 0, 0, 0, 0, warp_type(0))
+
+    @wp.func
+    def volatile_read_4d(ptr: wp.array(dtype=warp_type), index: int):
+        value = wp.atomic_exch(ptr, index, warp_type(0))
+        wp.atomic_exch(ptr, index, value)
+        return value
+
+    def test_spinlock_counter_4d(counter: wp.array(dtype=warp_type), lock: wp.array4d(dtype=warp_type)):
+        # Try to acquire the lock
+        spinlock_acquire_4d(lock)
+
+        # Critical section - increment counter
+        # counter[0] = counter[0] + 1 # This gives wrong results - counter should be marked as volatile
+
+        # Work around since warp arrays cannot be marked as volatile
+        value = volatile_read_4d(counter, 0)
+        counter[0] = value + warp_type(1)
+
+        # Release the lock
+        spinlock_release_4d(lock)
+
+    kernel = getkernel(test_spinlock_counter_4d, suffix=dtype.__name__)
+
+    if register_kernels:
+        return
+
+    wp.launch(kernel, dim=n, inputs=[counter, lock], device=device)
 
     # Verify counter reached n
     counter_np = counter.numpy()
-    expected = np.array([n], dtype=numpy_type)
+    expected = np.array([n], dtype=dtype)
 
     if not np.array_equal(counter_np, expected):
         print(f"Counter mismatch: expected {expected}, got {counter_np}")
@@ -244,54 +289,22 @@ class TestAtomicCAS(unittest.TestCase):
 
 
 # Test all supported types
-test_types = [
-    (wp.int32, np.int32),
-    (wp.uint32, np.uint32),
-    (wp.int64, np.int64),
-    (wp.uint64, np.uint64),
-    (wp.float32, np.float32),
-    (wp.float64, np.float64),
-]
-
-for warp_type, numpy_type in test_types:
-    type_name = warp_type.__name__
-    add_function_test(
-        TestAtomicCAS,
-        f"test_cas_{type_name}",
-        test_atomic_cas,
-        devices=devices,
-        warp_type=warp_type,
-        numpy_type=numpy_type,
-    )
+np_test_types = (np.int32, np.uint32, np.int64, np.uint64, np.float32, np.float64)
 
+for dtype in np_test_types:
+    type_name = dtype.__name__
+    add_function_test_register_kernel(
+        TestAtomicCAS, f"test_cas_{type_name}", test_atomic_cas, devices=devices, dtype=dtype
+    )
     # Add 2D test for each type
-    add_function_test(
-        TestAtomicCAS,
-        f"test_cas_2d_{type_name}",
-        test_atomic_cas_2d,
-        devices=devices,
-        warp_type=warp_type,
-        numpy_type=numpy_type,
+    add_function_test_register_kernel(
+        TestAtomicCAS, f"test_cas_2d_{type_name}", test_atomic_cas_2d, devices=devices, dtype=dtype
     )
-
-    # Add 3D test for each type
-    add_function_test(
-        TestAtomicCAS,
-        f"test_cas_3d_{type_name}",
-        test_atomic_cas_3d,
-        devices=devices,
-        warp_type=warp_type,
-        numpy_type=numpy_type,
+    add_function_test_register_kernel(
+        TestAtomicCAS, f"test_cas_3d_{type_name}", test_atomic_cas_3d, devices=devices, dtype=dtype
     )
-
-    # Add 4D test for each type
-    add_function_test(
-        TestAtomicCAS,
-        f"test_cas_4d_{type_name}",
-        test_atomic_cas_4d,
-        devices=devices,
-        warp_type=warp_type,
-        numpy_type=numpy_type,
+    add_function_test_register_kernel(
+        TestAtomicCAS, f"test_cas_4d_{type_name}", test_atomic_cas_4d, devices=devices, dtype=dtype
     )
 
 if __name__ == "__main__":

warp/tests/test_mat.py

@@ -2225,6 +2225,27 @@ def test_mat_array_sub_inplace(test, device):
     assert_np_equal(x.grad.numpy(), np.array([[[-1.0, -1.0], [-1.0, -1.0]]], dtype=float))
 
 
+@wp.kernel
+def scalar_mat_div(x: wp.array(dtype=wp.mat22), y: wp.array(dtype=wp.mat22)):
+    i = wp.tid()
+    y[i] = 1.0 / x[i]
+
+
+def test_scalar_mat_div(test, device):
+    x = wp.array((wp.mat22(1.0, 2.0, 4.0, 8.0),), dtype=wp.mat22, requires_grad=True, device=device)
+    y = wp.ones(1, dtype=wp.mat22, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(scalar_mat_div, 1, inputs=(x,), outputs=(y,), device=device)
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.array((((1.0, 0.5), (0.25, 0.125)),), dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array((((-1.0, -0.25), (-0.0625, -0.015625)),), dtype=float))
+
+
 devices = get_test_devices()
 
 
@@ -2356,6 +2377,7 @@ add_function_test(TestMat, "test_mat_add_inplace", test_mat_add_inplace, devices
 add_function_test(TestMat, "test_mat_sub_inplace", test_mat_sub_inplace, devices=devices)
 add_function_test(TestMat, "test_mat_array_add_inplace", test_mat_array_add_inplace, devices=devices)
 add_function_test(TestMat, "test_mat_array_sub_inplace", test_mat_array_sub_inplace, devices=devices)
+add_function_test(TestMat, "test_scalar_mat_div", test_scalar_mat_div, devices=devices)
 
 
 if __name__ == "__main__":

warp/tests/test_quat.py

@@ -2372,6 +2372,27 @@ def test_quat_array_sub_inplace(test, device):
     assert_np_equal(x.grad.numpy(), np.array([[-1.0, -1.0, -1.0, -1.0]], dtype=float))
 
 
+@wp.kernel
+def scalar_quat_div(x: wp.array(dtype=wp.quat), y: wp.array(dtype=wp.quat)):
+    i = wp.tid()
+    y[i] = 1.0 / x[i]
+
+
+def test_scalar_quat_div(test, device):
+    x = wp.array((wp.quat(1.0, 2.0, 4.0, 8.0),), dtype=wp.quat, requires_grad=True, device=device)
+    y = wp.ones(1, dtype=wp.quat, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(scalar_quat_div, 1, inputs=(x,), outputs=(y,), device=device)
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.array(((1.0, 0.5, 0.25, 0.125),), dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array(((-1.0, -0.25, -0.0625, -0.015625),), dtype=float))
+
+
 devices = get_test_devices()
 
 
@@ -2483,6 +2504,7 @@ add_function_test(TestQuat, "test_quat_add_inplace", test_quat_add_inplace, devi
 add_function_test(TestQuat, "test_quat_sub_inplace", test_quat_sub_inplace, devices=devices)
 add_function_test(TestQuat, "test_quat_array_add_inplace", test_quat_array_add_inplace, devices=devices)
 add_function_test(TestQuat, "test_quat_array_sub_inplace", test_quat_array_sub_inplace, devices=devices)
+add_function_test(TestQuat, "test_scalar_quat_div", test_scalar_quat_div, devices=devices)
 
 
 if __name__ == "__main__":

warp/tests/test_sparse.py

@@ -140,6 +140,32 @@ def test_bsr_from_triplets(test, device):
         bsr_set_from_triplets(bsr, rows, cols, vals)
 
 
+def test_bsr_from_triplets_prune_numerical_zeros(test, device):
+    rows = wp.array([1, 0, 2, 3], dtype=int)
+    cols = wp.array([0, 1, 2, 3], dtype=int)
+    vals = wp.zeros(len(rows), dtype=float)
+
+    A = bsr_from_triplets(
+        rows_of_blocks=12,  # Number of rows of blocks
+        cols_of_blocks=12,  # Number of columns of blocks
+        rows=rows,  # Row indices
+        columns=cols,  # Column indices
+        values=vals,  # Block values
+        prune_numerical_zeros=False,
+    )
+    assert A.nnz_sync() == 4
+
+    A = bsr_from_triplets(
+        rows_of_blocks=12,  # Number of rows of blocks
+        cols_of_blocks=12,  # Number of columns of blocks
+        rows=rows,  # Row indices
+        columns=cols,  # Column indices
+        values=vals,  # Block values
+        prune_numerical_zeros=True,
+    )
+    assert A.nnz_sync() == 0
+
+
 def test_bsr_from_triplets_gradient(test, device):
     rng = np.random.default_rng(123)
 
@@ -604,6 +630,12 @@ class TestSparse(unittest.TestCase):
 
 add_function_test(TestSparse, "test_csr_from_triplets", test_csr_from_triplets, devices=devices)
 add_function_test(TestSparse, "test_bsr_from_triplets", test_bsr_from_triplets, devices=devices)
+add_function_test(
+    TestSparse,
+    "test_bsr_from_triplets_prune_numerical_zeros",
+    test_bsr_from_triplets_prune_numerical_zeros,
+    devices=devices,
+)
 add_function_test(TestSparse, "test_bsr_get_diag", test_bsr_get_set_diag, devices=devices)
 add_function_test(TestSparse, "test_bsr_split_merge", test_bsr_split_merge, devices=devices)
 add_function_test(TestSparse, "test_bsr_assign_masked", test_bsr_assign_masked, devices=devices)

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