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

warp/tests/test_quat_assign_copy.py

@@ -0,0 +1,145 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+def setUpModule():
+    wp.config.enable_vector_component_overwrites = True
+
+
+def tearDownModule():
+    wp.config.enable_vector_component_overwrites = False
+
+
+@wp.kernel
+def quat_assign_subscript(x: wp.array(dtype=float), y: wp.array(dtype=wp.quat)):
+    i = wp.tid()
+
+    a = wp.quat()
+    a[0] = 1.0 * x[i]
+    a[1] = 2.0 * x[i]
+    a[2] = 3.0 * x[i]
+    a[3] = 4.0 * x[i]
+    y[i] = a
+
+
+@wp.kernel
+def quat_assign_attribute(x: wp.array(dtype=float), y: wp.array(dtype=wp.quat)):
+    i = wp.tid()
+
+    a = wp.quat()
+    a.x = 1.0 * x[i]
+    a.y = 2.0 * x[i]
+    a.z = 3.0 * x[i]
+    a.w = 4.0 * x[i]
+    y[i] = a
+
+
+def test_quat_assign(test, device):
+    def run(kernel):
+        x = wp.ones(1, dtype=float, requires_grad=True, device=device)
+        y = wp.zeros(1, dtype=wp.quat, requires_grad=True, device=device)
+
+        tape = wp.Tape()
+        with tape:
+            wp.launch(kernel, 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]], dtype=float))
+        assert_np_equal(x.grad.numpy(), np.array([10.0], dtype=float))
+
+    run(quat_assign_subscript)
+    run(quat_assign_attribute)
+
+
+def test_quat_assign_copy(test, device):
+    @wp.kernel(module="unique")
+    def quat_assign_overwrite(x: wp.array(dtype=wp.quat), y: wp.array(dtype=wp.quat)):
+        tid = wp.tid()
+
+        a = wp.quat()
+        b = x[tid]
+        a = b
+        a[1] = 3.0
+
+        y[tid] = a
+
+    x = wp.ones(1, dtype=wp.quat, device=device, requires_grad=True)
+    y = wp.zeros(1, dtype=wp.quat, device=device, requires_grad=True)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(quat_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]], dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array([[1.0, 0.0, 1.0, 1.0]], dtype=float))
+
+
+def test_quat_slicing_assign_backward(test, device):
+    @wp.kernel(module="unique")
+    def kernel(arr_x: wp.array(dtype=wp.vec2), arr_y: wp.array(dtype=wp.quat)):
+        i = wp.tid()
+
+        x = arr_x[i]
+        y = arr_y[i]
+
+        y[:2] = x
+        y[1:-1] += x[:2]
+        y[3:1:-1] -= x[0:]
+
+        arr_y[i] = y
+
+    x = wp.ones(1, dtype=wp.vec2, requires_grad=True, device=device)
+    y = wp.zeros(1, dtype=wp.quat, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(kernel, 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, 0.0, -1.0),), dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array(((1.0, 1.0),), dtype=float))
+
+
+devices = get_test_devices()
+
+
+class TestQuatAssignCopy(unittest.TestCase):
+    pass
+
+
+add_function_test(TestQuatAssignCopy, "test_quat_assign", test_quat_assign, devices=devices)
+add_function_test(TestQuatAssignCopy, "test_quat_assign_copy", test_quat_assign_copy, devices=devices)
+add_function_test(
+    TestQuatAssignCopy, "test_quat_slicing_assign_backward", test_quat_slicing_assign_backward, devices=devices
+)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_reload.py

@@ -35,7 +35,8 @@ def reload_module(module):
     # Clearing the .pyc file associated with a module is a necessary workaround
     # for `importlib.reload` to work as expected when run from within Kit.
     cache_file = importlib.util.cache_from_source(module.__file__)
-    os.remove(cache_file)
+    if os.path.exists(cache_file):
+        os.remove(cache_file)
     importlib.reload(module)
 
 

warp/tests/test_sparse.py

@@ -568,6 +568,7 @@ def make_test_bsr_mv(block_shape, scalar_type):
         work_buffer = wp.empty_like(y)
         for alpha, beta in zip(alphas, betas):
             ref = alpha * _bsr_to_dense(A) @ x.numpy().flatten() + beta * y.numpy().flatten()
+
             if beta == 0.0:
                 y = A @ x
             else:
@@ -604,7 +605,67 @@ def make_test_bsr_mv(block_shape, scalar_type):
     return test_bsr_mv
 
 
+def make_test_bsr_multiply_deep(block_shape, scalar_type):
+    def test_bsr_multiply_deep(test, device):
+        """Test BSR matrix multiplication with deep matrices (many columns > 256)"""
+        rng = np.random.default_rng(123)
+
+        # Generate a dense matrix with few rows and many columns (> 256)
+        nrow = (4 + block_shape[0] - 1) // block_shape[0]
+        ncol = (600 + block_shape[1] - 1) // block_shape[1]
+
+        # Create a dense "sparse" matrix
+        values = rng.random(size=(nrow * ncol, block_shape[0], block_shape[1]))
+        rows, cols = np.meshgrid(np.arange(nrow), np.arange(ncol))
+
+        # Convert to warp arrays
+        rows = wp.array(rows.flatten(), dtype=int, device=device)
+        cols = wp.array(cols.flatten(), dtype=int, device=device)
+        vals = wp.array(values, dtype=scalar_type, device=device)
+
+        # Convert to BSR using bsr_from_triplets
+        A = bsr_from_triplets(nrow, ncol, rows, cols, vals)
+
+        # Get dense representation for numpy reference
+        A_dense = _bsr_to_dense(A)
+
+        # Multiply with itself transpose using bsr_mm
+        # A @ A.T should result in a nrow x nrow matrix
+        At = bsr_transposed(A)
+
+        result = bsr_mm(A, At)
+
+        # Check that the result is correct against numpy reference
+        result_dense = _bsr_to_dense(result)
+        ref_dense = A_dense @ A_dense.T
+
+        assert_np_equal(result_dense, ref_dense, 0.0001)
+
+        # Additional test: multiply A.T @ A (should be ncol x ncol)
+        result2 = bsr_mm(At, A)
+        result2_dense = _bsr_to_dense(result2)
+        ref2_dense = A_dense.T @ A_dense
+
+        assert_np_equal(result2_dense, ref2_dense, 0.0001)
+
+        # Test matrix vector products
+        x = wp.array(rng.random(size=A.shape[1]), dtype=A.scalar_type, device=device)
+        y = wp.array(rng.random(size=A.shape[0]), dtype=A.scalar_type, device=device)
+        bsr_mv(A, x, y)
+        res = y.numpy().flatten()
+        ref = A_dense @ x.numpy().flatten()
+        assert_np_equal(res, ref, 0.0001 * block_shape[1])
+
+        bsr_mv(A, y, x, transpose=True)
+        res = x.numpy().flatten()
+        ref = A_dense.T @ y.numpy().flatten()
+        assert_np_equal(res, ref, 0.0001 * block_shape[1])
+
+    return test_bsr_multiply_deep
+
+
 devices = get_test_devices()
+cuda_test_devices = get_selected_cuda_test_devices()
 
 
 class TestSparse(unittest.TestCase):
@@ -653,6 +714,16 @@ add_function_test(TestSparse, "test_csr_mm", make_test_bsr_mm((1, 1), wp.float32
 add_function_test(TestSparse, "test_bsr_mm_1_3", make_test_bsr_mm((1, 3), wp.float32), devices=devices)
 add_function_test(TestSparse, "test_bsr_mm_3_3", make_test_bsr_mm((3, 3), wp.float64), devices=devices)
 
+add_function_test(
+    TestSparse, "test_bsr_multiply_deep_2_2", make_test_bsr_multiply_deep((2, 2), wp.float64), devices=devices
+)
+add_function_test(
+    TestSparse,
+    "test_bsr_multiply_deep_30_30",
+    make_test_bsr_multiply_deep((30, 30), wp.float32),
+    devices=cuda_test_devices,
+)
+
 add_function_test(TestSparse, "test_csr_mv", make_test_bsr_mv((1, 1), wp.float32), devices=devices)
 add_function_test(TestSparse, "test_bsr_mv_1_3", make_test_bsr_mv((1, 3), wp.float32), devices=devices)
 add_function_test(TestSparse, "test_bsr_mv_3_3", make_test_bsr_mv((3, 3), wp.float64), devices=devices)

warp/tests/test_spatial.py

@@ -2217,39 +2217,6 @@ def test_transform_assign(test, device):
     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()
@@ -2510,6 +2477,141 @@ def test_transform_array_sub_inplace(test, device):
     assert_np_equal(x.grad.numpy(), np.array([[-1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0]], dtype=float))
 
 
+def test_transform_indexing_assign(test, device):
+    @wp.func
+    def fn():
+        t = wp.transform(p=wp.vec3(1.0, 2.0, 3.0), q=wp.quat(4.0, 5.0, 6.0, 7.0))
+
+        t[0] = 123.0
+        t[3] *= 2.0
+
+        wp.expect_eq(t[0], 123.0)
+        wp.expect_eq(t[1], 2.0)
+        wp.expect_eq(t[2], 3.0)
+        wp.expect_eq(t[3], 8.0)
+        wp.expect_eq(t[4], 5.0)
+        wp.expect_eq(t[5], 6.0)
+        wp.expect_eq(t[6], 7.0)
+
+        t[-1] = 123.0
+        t[-5] *= 2.0
+
+        wp.expect_eq(t[0], 123.0)
+        wp.expect_eq(t[1], 2.0)
+        wp.expect_eq(t[2], 6.0)
+        wp.expect_eq(t[3], 8.0)
+        wp.expect_eq(t[4], 5.0)
+        wp.expect_eq(t[5], 6.0)
+        wp.expect_eq(t[6], 123.0)
+
+    @wp.kernel(module="unique")
+    def kernel():
+        fn()
+
+    wp.launch(kernel, 1, device=device)
+    wp.synchronize()
+    fn()
+
+
+def test_transform_slicing_assign(test, device):
+    vec0 = wp.vec(0, float)
+    vec1 = wp.vec(1, float)
+    vec2 = wp.vec(2, float)
+    vec3 = wp.vec(3, float)
+    vec4 = wp.vec(4, float)
+    vec5 = wp.vec(5, float)
+    vec6 = wp.vec(6, float)
+    vec7 = wp.vec(7, float)
+
+    @wp.func
+    def fn():
+        t = wp.transform(p=wp.vec3(1.0, 2.0, 3.0), q=wp.quat(4.0, 5.0, 6.0, 7.0))
+
+        wp.expect_eq(t[:] == vec7(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0), True)
+        wp.expect_eq(t[-123:123] == vec7(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0), True)
+        wp.expect_eq(t[123:] == vec0(), True)
+        wp.expect_eq(t[:-123] == vec0(), True)
+        wp.expect_eq(t[::123] == vec1(1.0), True)
+
+        wp.expect_eq(t[1:] == vec6(2.0, 3.0, 4.0, 5.0, 6.0, 7.0), True)
+        wp.expect_eq(t[-2:] == vec2(6.0, 7.0), True)
+        wp.expect_eq(t[:2] == vec2(1.0, 2.0), True)
+        wp.expect_eq(t[:-1] == vec6(1.0, 2.0, 3.0, 4.0, 5.0, 6.0), True)
+        wp.expect_eq(t[::2] == vec4(1.0, 3.0, 5.0, 7.0), True)
+        wp.expect_eq(t[1::2] == vec3(2.0, 4.0, 6.0), True)
+        wp.expect_eq(t[::-1] == vec7(7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0), True)
+        wp.expect_eq(t[::-2] == vec4(7.0, 5.0, 3.0, 1.0), True)
+        wp.expect_eq(t[1::-2] == vec1(2.0), True)
+
+        t[1:] = vec6(8.0, 9.0, 10.0, 11.0, 12.0, 13.0)
+        wp.expect_eq(t == wp.transform(p=wp.vec3(1.0, 8.0, 9.0), q=wp.quat(10.0, 11.0, 12.0, 13.0)), True)
+
+        t[-2:] = vec2(14.0, 15.0)
+        wp.expect_eq(t == wp.transform(p=wp.vec3(1.0, 8.0, 9.0), q=wp.quat(10.0, 11.0, 14.0, 15.0)), True)
+
+        t[:2] = vec2(16.0, 17.0)
+        wp.expect_eq(t == wp.transform(p=wp.vec3(16.0, 17.0, 9.0), q=wp.quat(10.0, 11.0, 14.0, 15.0)), True)
+
+        t[:-1] = vec6(18.0, 19.0, 20.0, 21.0, 22.0, 23.0)
+        wp.expect_eq(t == wp.transform(p=wp.vec3(18.0, 19.0, 20.0), q=wp.quat(21.0, 22.0, 23.0, 15.0)), True)
+
+        t[::2] = vec4(24.0, 25.0, 26.0, 27.0)
+        wp.expect_eq(t == wp.transform(p=wp.vec3(24.0, 19.0, 25.0), q=wp.quat(21.0, 26.0, 23.0, 27.0)), True)
+
+        t[1::2] = vec3(28.0, 29.0, 30.0)
+        wp.expect_eq(t == wp.transform(p=wp.vec3(24.0, 28.0, 25.0), q=wp.quat(29.0, 26.0, 30.0, 27.0)), True)
+
+        t[::-1] = vec7(31.0, 32.0, 33.0, 34.0, 35.0, 36.0, 37.0)
+        wp.expect_eq(t == wp.transform(p=wp.vec3(37.0, 36.0, 35.0), q=wp.quat(34.0, 33.0, 32.0, 31.0)), True)
+
+        t[::-2] = vec4(38.0, 39.0, 40.0, 41.0)
+        wp.expect_eq(t == wp.transform(p=wp.vec3(41.0, 36.0, 40.0), q=wp.quat(34.0, 39.0, 32.0, 38.0)), True)
+
+        t[1::-2] = vec1(42.0)
+        wp.expect_eq(t == wp.transform(p=wp.vec3(41.0, 42.0, 40.0), q=wp.quat(34.0, 39.0, 32.0, 38.0)), True)
+
+        t[1:] += vec6(43.0, 44.0, 45.0, 46.0, 47.0, 48.0)
+        wp.expect_eq(t == wp.transform(p=wp.vec3(41.0, 85.0, 84.0), q=wp.quat(79.0, 85.0, 79.0, 86.0)), True)
+
+        t[:-1] -= vec6(49.0, 50.0, 51.0, 52.0, 53.0, 54.0)
+        wp.expect_eq(t == wp.transform(p=wp.vec3(-8.0, 35.0, 33.0), q=wp.quat(27.0, 32.0, 25.0, 86.0)), True)
+
+    @wp.kernel(module="unique")
+    def kernel():
+        fn()
+
+    wp.launch(kernel, 1, device=device)
+    wp.synchronize()
+    fn()
+
+
+def test_transform_slicing_assign_backward(test, device):
+    @wp.kernel(module="unique")
+    def kernel(arr_x: wp.array(dtype=wp.vec2), arr_y: wp.array(dtype=wp.transform)):
+        i = wp.tid()
+
+        y = arr_y[i]
+
+        y[:2] = arr_x[i]
+        y[1:-4] += arr_x[i][:2]
+        y[3:1:-1] -= arr_x[i][0:]
+
+        arr_y[i] = y
+
+    x = wp.ones(1, dtype=wp.vec2, requires_grad=True, device=device)
+    y = wp.zeros(1, dtype=wp.transform, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(kernel, 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, 0.0, -1.0, 0.0, 0.0, 0.0),), dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array(((1.0, 1.0),), dtype=float))
+
+
 devices = get_test_devices()
 
 
@@ -2721,13 +2823,17 @@ add_function_test(
 )
 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)
+add_function_test(TestSpatial, "test_transform_indexing_assign", test_transform_indexing_assign, devices=devices)
+add_function_test(TestSpatial, "test_transform_slicing_assign", test_transform_slicing_assign, devices=devices)
+add_function_test(
+    TestSpatial, "test_transform_slicing_assign_backward", test_transform_slicing_assign_backward, devices=devices
+)
 
 
 if __name__ == "__main__":

warp/tests/test_spatial_assign_copy.py

@@ -0,0 +1,160 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+def setUpModule():
+    wp.config.enable_vector_component_overwrites = True
+
+
+def tearDownModule():
+    wp.config.enable_vector_component_overwrites = False
+
+
+@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):
+    @wp.kernel(module="unique")
+    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))
+
+
+def test_transform_slicing_assign_backward(test, device):
+    @wp.kernel(module="unique")
+    def kernel(arr_x: wp.array(dtype=wp.vec2), arr_y: wp.array(dtype=wp.transform)):
+        i = wp.tid()
+
+        x = arr_x[i]
+        y = arr_y[i]
+
+        y[:2] = x
+        y[1:-4] += x[:2]
+        y[3:1:-1] -= x[0:]
+
+        arr_y[i] = y
+
+    x = wp.ones(1, dtype=wp.vec2, requires_grad=True, device=device)
+    y = wp.zeros(1, dtype=wp.transform, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(kernel, 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, 0.0, -1.0, 0.0, 0.0, 0.0),), dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array(((1.0, 1.0),), dtype=float))
+
+
+devices = get_test_devices()
+
+
+class TestSpatialAssignCopy(unittest.TestCase):
+    pass
+
+
+add_function_test(TestSpatialAssignCopy, "test_transform_assign", test_transform_assign, devices=devices)
+add_function_test(TestSpatialAssignCopy, "test_transform_assign_copy", test_transform_assign_copy, devices=devices)
+add_function_test(
+    TestSpatialAssignCopy,
+    "test_transform_slicing_assign_backward",
+    test_transform_slicing_assign_backward,
+    devices=devices,
+)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

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)