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

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

@@ -235,6 +235,44 @@ def test_nested_struct(test, device):
     )
 
 
+@wp.struct
+class MatStruct:
+    m: wp.mat44
+
+
+@wp.kernel
+def kernel_nested_mat(out: wp.array(dtype=wp.mat44)):
+    s = MatStruct()
+    m = wp.mat44()
+
+    s.m[1, 2] = 3.0
+
+    out[0] = s.m
+
+
+def test_nested_mat(test, device):
+    m = wp.array([wp.mat44()], dtype=wp.mat44, device=device)
+    wp.launch(kernel_nested_mat, dim=1, outputs=[m], device=device)
+    wp.synchronize()
+
+    out = m.numpy()
+    assert_np_equal(out[0][1, 2], 3.0)
+
+
+def test_assign_view(test, device):
+    @wp.kernel
+    def kernel_assign_view(out: wp.array2d(dtype=wp.mat44)):
+        out[0][2, 2] = 6.0
+
+    m = wp.array([[wp.mat44()]], dtype=wp.mat44, device=device)
+
+    with test.assertRaisesRegex(
+        wp.codegen.WarpCodegenError,
+        r"Incorrect number of indices specified for array indexing",
+    ):
+        wp.launch(kernel_assign_view, dim=[1, 1], outputs=[m], device=device)
+
+
 def test_struct_attribute_error(test, device):
     @wp.kernel
     def kernel(foo: Foo):
@@ -751,9 +789,6 @@ def test_struct_array_gc_requires_grad_toggle(test, device):
     tape.backward(loss=loss_wp)
 
 
-devices = get_test_devices()
-
-
 class TestStruct(unittest.TestCase):
     # check structs default initialized in Python correctly
     def test_struct_default_attributes_python(self):
@@ -772,6 +807,9 @@ class TestStruct(unittest.TestCase):
         np.testing.assert_equal(arr.numpy().tolist(), expected)
 
 
+devices = get_test_devices()
+
+
 add_function_test(TestStruct, "test_step", test_step, devices=devices)
 add_function_test(TestStruct, "test_step_grad", test_step_grad, devices=devices)
 add_kernel_test(TestStruct, kernel=test_empty, name="test_empty", dim=1, inputs=[Empty()], devices=devices)
@@ -785,6 +823,8 @@ add_kernel_test(
 )
 add_kernel_test(TestStruct, kernel=test_return, name="test_return", dim=1, inputs=[], devices=devices)
 add_function_test(TestStruct, "test_nested_struct", test_nested_struct, devices=devices)
+add_function_test(TestStruct, "test_nested_mat", test_nested_mat, devices=devices)
+add_function_test(TestStruct, "test_assign_view", test_assign_view, devices=devices)
 add_function_test(TestStruct, "test_struct_attribute_error", test_struct_attribute_error, devices=devices)
 add_function_test(TestStruct, "test_struct_inheritance_error", test_struct_inheritance_error, devices=devices)
 add_function_test(TestStruct, "test_nested_array_struct", test_nested_array_struct, 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)
 
 

warp/tests/test_types.py

@@ -267,11 +267,6 @@ class TestTypes(unittest.TestCase):
         with self.assertRaisesRegex(TypeError, r"Expected to assign a `int32` value but got `str` instead"):
             v1[0] = "123.0"
 
-        with self.assertRaisesRegex(
-            TypeError, r"Expected to assign a slice from a sequence of values but got `int` instead"
-        ):
-            v1[:] = 123
-
         with self.assertRaisesRegex(
             TypeError, r"Expected to assign a slice from a sequence of `int32` values but got `vec3i` instead"
         ):
@@ -483,26 +478,11 @@ class TestTypes(unittest.TestCase):
         with self.assertRaisesRegex(TypeError, r"Expected to assign a `float16` value but got `str` instead"):
             m[0][0] = "123.0"
 
-        with self.assertRaisesRegex(
-            TypeError, r"Expected to assign a slice from a sequence of values but got `int` instead"
-        ):
-            m[0] = 123
-
         with self.assertRaisesRegex(
             TypeError, r"Expected to assign a slice from a sequence of `float16` values but got `mat22h` instead"
         ):
             m[0] = (m,)
 
-        with self.assertRaisesRegex(
-            KeyError, r"Slices are not supported when indexing matrices using the `m\[start:end\]` notation"
-        ):
-            m[:] = 123
-
-        with self.assertRaisesRegex(
-            KeyError, r"Slices are not supported when indexing matrices using the `m\[i, j\]` notation"
-        ):
-            m[0, :1] = (123,)
-
         with self.assertRaisesRegex(ValueError, r"Can only assign sequence of same size"):
             m[0][:1] = (1, 2)