warp-lang 1.8.1__py3-none-win_amd64.whl → 1.9.1__py3-none-win_amd64.whl

warp/tests/test_tuple.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import sys
 import unittest
 from typing import Any, Tuple
 
@@ -208,6 +209,92 @@ def test_loop_variadic_ellipsis():
     wp.expect_eq(res, 22)
 
 
+# Test for Python 3.10 tuple type compatibility issue
+# Only define these functions on Python 3.9+ where lowercase tuple is supported
+if sys.version_info >= (3, 9):
+
+    @wp.func
+    def complex_tuple_function(scale: float, offset: wp.vec3) -> tuple[float, wp.vec3f, wp.vec3f]:
+        """
+        Function that returns a complex tuple with mixed types.
+        This specifically tests the tuple[float, wp.vec3f, wp.vec3f] case
+        that was problematic on Python 3.10.
+        """
+        # Create some computed values
+        scaled_value = scale * 2.5
+        position = wp.vec3f(offset.x + 1.0, offset.y + 2.0, offset.z + 3.0)
+        velocity = wp.vec3f(scale * 0.1, scale * 0.2, scale * 0.3)
+
+        return (scaled_value, position, velocity)
+
+    @wp.func
+    def mixed_types_tuple_function() -> tuple[wp.vec3f, wp.vec3f, float, wp.mat33f]:
+        """
+        Function returning mixed types in a tuple.
+        Tests tuple[vec3f, vec3f, float, mat33f] type annotation.
+        """
+        return (
+            wp.vec3f(1.0, 2.0, 3.0),
+            wp.vec3f(4.0, 5.0, 6.0),
+            42.0,
+            wp.mat33f(1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0),
+        )
+
+    @wp.func
+    def homogeneous_tuple_function() -> tuple[wp.vec3f, wp.vec3f, wp.vec3f]:
+        """
+        Function returning fixed-size homogeneous tuple.
+        Tests tuple[wp.vec3f, wp.vec3f, wp.vec3f] type annotation.
+        """
+        return (wp.vec3f(1.0, 2.0, 3.0), wp.vec3f(4.0, 5.0, 6.0), wp.vec3f(7.0, 8.0, 9.0))
+
+    @wp.kernel
+    def test_complex_tuple_functions():
+        """
+        Kernel that tests complex tuple return types that were problematic on Python 3.10.
+        """
+        # Test the main problematic case: tuple[float, wp.vec3f, wp.vec3f]
+        result1 = complex_tuple_function(4.0, wp.vec3(10.0, 20.0, 30.0))
+
+        # Unpack and verify
+        scale_result, pos_result, vel_result = result1
+        wp.expect_near(scale_result, 10.0)  # 4.0 * 2.5
+        wp.expect_eq(pos_result, wp.vec3f(11.0, 22.0, 33.0))
+        wp.expect_eq(vel_result, wp.vec3f(0.4, 0.8, 1.2))
+
+        # Test access by index
+        wp.expect_near(result1[0], 10.0)
+        wp.expect_eq(result1[1], wp.vec3f(11.0, 22.0, 33.0))
+        wp.expect_eq(result1[2], wp.vec3f(0.4, 0.8, 1.2))
+
+        # Test more complex tuple: tuple[vec3f, vec3f, float, mat33f]
+        mixed_result = mixed_types_tuple_function()
+        result_pos, result_vel, result_energy, result_transform = mixed_result
+
+        # Verify known values
+        wp.expect_eq(result_pos, wp.vec3f(1.0, 2.0, 3.0))
+        wp.expect_eq(result_vel, wp.vec3f(4.0, 5.0, 6.0))
+        wp.expect_eq(result_energy, 42.0)
+
+        # Verify transform matrix is identity
+        wp.expect_eq(result_transform[0, 0], 1.0)
+        wp.expect_eq(result_transform[1, 1], 1.0)
+        wp.expect_eq(result_transform[2, 2], 1.0)
+
+        # Test fixed-size homogeneous tuple: tuple[wp.vec3f, wp.vec3f, wp.vec3f]
+        homo_result = homogeneous_tuple_function()
+        wp.expect_eq(len(homo_result), 3)
+        wp.expect_eq(homo_result[0], wp.vec3f(1.0, 2.0, 3.0))
+        wp.expect_eq(homo_result[1], wp.vec3f(4.0, 5.0, 6.0))
+        wp.expect_eq(homo_result[2], wp.vec3f(7.0, 8.0, 9.0))
+
+        # Test unpacking
+        vec1, vec2, vec3 = homo_result
+        wp.expect_eq(vec1, wp.vec3f(1.0, 2.0, 3.0))
+        wp.expect_eq(vec2, wp.vec3f(4.0, 5.0, 6.0))
+        wp.expect_eq(vec3, wp.vec3f(7.0, 8.0, 9.0))
+
+
 devices = get_test_devices()
 
 
@@ -258,6 +345,15 @@ add_kernel_test(
     dim=1,
     devices=devices,
 )
+# Only register the test for lowercase tuple syntax on Python 3.9+
+if sys.version_info >= (3, 9):
+    add_kernel_test(
+        TestTuple,
+        name="test_complex_tuple_functions",
+        kernel=test_complex_tuple_functions,
+        dim=1,
+        devices=devices,
+    )
 
 
 if __name__ == "__main__":

warp/tests/test_types.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import sys
 import unittest
 
 from warp.tests.unittest_utils import *
@@ -267,11 +268,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 +479,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)
 
@@ -547,6 +528,66 @@ class TestTypes(unittest.TestCase):
         test_conversions(wp.uint64, np.uint64)
         test_conversions(wp.bool, np.bool_)
 
+    # Only define this test method on Python 3.9+ where lowercase tuple is supported
+    if sys.version_info >= (3, 9):
+
+        def test_tuple_type_code_generation(self):
+            """Test that tuple type annotations generate correct type codes, especially on Python 3.10."""
+            import sys
+            import types
+            from typing import get_origin
+
+            # Test basic tuple types
+            tuple_float_float = tuple[float, float]
+            result = wp.types.get_type_code(tuple_float_float)
+            self.assertEqual(result, "tpl2f4f4", "tuple[float, float] should generate 'tpl2f4f4'")
+
+            # Test tuple with Warp vector types - the problematic case from Python 3.10
+            tuple_mixed = tuple[float, wp.vec3f, wp.vec3f]
+            result = wp.types.get_type_code(tuple_mixed)
+            self.assertEqual(result, "tpl3f4v3f4v3f4", "tuple[float, vec3f, vec3f] should generate 'tpl3f4v3f4v3f4'")
+
+            # Test homogeneous tuple with ellipsis
+            tuple_homogeneous = tuple[wp.vec3f, ...]
+            result = wp.types.get_type_code(tuple_homogeneous)
+            self.assertEqual(result, "tpl2v3f4?", "tuple[vec3f, ...] should generate 'tpl2v3f4?'")
+
+            # Test single element tuple
+            tuple_single = tuple[wp.int32]
+            result = wp.types.get_type_code(tuple_single)
+            self.assertEqual(result, "tpl1i4", "tuple[int32] should generate 'tpl1i4'")
+
+            # Test tuple with multiple Warp types
+            tuple_multi_warp = tuple[wp.vec3f, wp.mat33f, wp.quatf]
+            result = wp.types.get_type_code(tuple_multi_warp)
+            self.assertEqual(
+                result, "tpl3v3f4m33f4qf4", "tuple[vec3f, mat33f, quatf] should generate 'tpl3v3f4m33f4qf4'"
+            )
+
+            # Verify the fix works on Python 3.9-3.10 specifically
+            if sys.version_info < (3, 11) and hasattr(types, "GenericAlias"):
+                # On Python 3.9-3.10, tuple[...] creates types.GenericAlias
+                self.assertIsInstance(
+                    tuple_mixed, types.GenericAlias, "On Python 3.9-3.10, tuple[...] should create types.GenericAlias"
+                )
+                self.assertIs(tuple_mixed.__origin__, tuple, "GenericAlias origin should be tuple")
+
+                # Verify our fix catches this case
+                self.assertEqual(get_origin(tuple_mixed), tuple, "get_origin should return tuple")
+            elif sys.version_info >= (3, 11):
+                # On Python 3.11+, the existing code path should handle it
+                self.assertEqual(get_origin(tuple_mixed), tuple, "get_origin should return tuple on Python 3.11+")
+
+            # Test that the fix doesn't break existing functionality
+            # Test with built-in Python types
+            tuple_builtin = tuple[int, str, bool]
+            try:
+                # This might fail because str and bool aren't Warp types, but it shouldn't crash
+                wp.types.get_type_code(tuple_builtin)
+            except TypeError as e:
+                # Expected to fail for non-Warp types, but should be a clean TypeError
+                self.assertIn("Unrecognized type", str(e))
+
 
 for dtype in wp.types.int_types:
     add_function_test(TestTypes, f"test_integers_{dtype.__name__}", test_integers, devices=devices, dtype=dtype)

warp/tests/test_vec.py

@@ -922,39 +922,6 @@ def test_vec_assign(test, device):
     run(vec_assign_attribute)
 
 
-def test_vec_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(module="unique")
-        def vec_assign_overwrite(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=wp.vec3)):
-            tid = wp.tid()
-
-            a = wp.vec3()
-            b = x[tid]
-            a = b
-            a[1] = 3.0
-
-            y[tid] = a
-
-        x = wp.ones(1, dtype=wp.vec3, device=device, requires_grad=True)
-        y = wp.zeros(1, dtype=wp.vec3, device=device, requires_grad=True)
-
-        tape = wp.Tape()
-        with tape:
-            wp.launch(vec_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]], dtype=float))
-        assert_np_equal(x.grad.numpy(), np.array([[1.0, 0.0, 1.0]], dtype=float))
-
-    finally:
-        wp.config.enable_vector_component_overwrites = saved_enable_vector_component_overwrites_setting
-
-
 @wp.kernel
 def vec_array_extract_subscript(x: wp.array2d(dtype=wp.vec3), y: wp.array2d(dtype=float)):
     i, j = wp.tid()
@@ -1189,6 +1156,181 @@ def test_scalar_vec_div(test, device):
     assert_np_equal(x.grad.numpy(), np.array(((-1.0, -0.25, -0.0625),), dtype=float))
 
 
+def test_vec_indexing_assign(test, device):
+    @wp.func
+    def fn():
+        v = wp.vec4(1.0, 2.0, 3.0, 4.0)
+
+        v[0] = 123.0
+        v[1] *= 2.0
+
+        wp.expect_eq(v[0], 123.0)
+        wp.expect_eq(v[1], 4.0)
+        wp.expect_eq(v[2], 3.0)
+        wp.expect_eq(v[3], 4.0)
+
+        v[-1] = 123.0
+        v[-2] *= 2.0
+
+        wp.expect_eq(v[-1], 123.0)
+        wp.expect_eq(v[-2], 6.0)
+        wp.expect_eq(v[-3], 4.0)
+        wp.expect_eq(v[-4], 123.0)
+
+    @wp.kernel(module="unique")
+    def kernel():
+        fn()
+
+    wp.launch(kernel, 1, device=device)
+    wp.synchronize()
+    fn()
+
+
+def test_vec_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)
+
+    @wp.func
+    def fn():
+        v = wp.vec4(1.0, 2.0, 3.0, 4.0)
+
+        wp.expect_eq(v[:] == vec4(1.0, 2.0, 3.0, 4.0), True)
+        wp.expect_eq(v[-123:123] == vec4(1.0, 2.0, 3.0, 4.0), True)
+        wp.expect_eq(v[123:] == vec0(), True)
+        wp.expect_eq(v[:-123] == vec0(), True)
+        wp.expect_eq(v[::123] == vec1(1.0), True)
+
+        wp.expect_eq(v[1:] == vec3(2.0, 3.0, 4.0), True)
+        wp.expect_eq(v[-2:] == vec2(3.0, 4.0), True)
+        wp.expect_eq(v[:2] == vec2(1.0, 2.0), True)
+        wp.expect_eq(v[:-1] == vec3(1.0, 2.0, 3.0), True)
+        wp.expect_eq(v[::2] == vec2(1.0, 3.0), True)
+        wp.expect_eq(v[1::2] == vec2(2.0, 4.0), True)
+        wp.expect_eq(v[::-1] == vec4(4.0, 3.0, 2.0, 1.0), True)
+        wp.expect_eq(v[::-2] == vec2(4.0, 2.0), True)
+        wp.expect_eq(v[1::-2] == vec1(2.0), True)
+
+        v[1:] = vec3(5.0, 6.0, 7.0)
+        wp.expect_eq(v == wp.vec4(1.0, 5.0, 6.0, 7.0), True)
+
+        v[-2:] = vec2(8.0, 9.0)
+        wp.expect_eq(v == wp.vec4(1.0, 5.0, 8.0, 9.0), True)
+
+        v[:2] = vec2(10.0, 11.0)
+        wp.expect_eq(v == wp.vec4(10.0, 11.0, 8.0, 9.0), True)
+
+        v[:-1] = vec3(12.0, 13.0, 14.0)
+        wp.expect_eq(v == wp.vec4(12.0, 13.0, 14.0, 9.0), True)
+
+        v[::2] = vec2(15.0, 16.0)
+        wp.expect_eq(v == wp.vec4(15.0, 13.0, 16.0, 9.0), True)
+
+        v[1::2] = vec2(17.0, 18.0)
+        wp.expect_eq(v == wp.vec4(15.0, 17.0, 16.0, 18.0), True)
+
+        v[::-1] = vec4(19.0, 20.0, 21.0, 22.0)
+        wp.expect_eq(v == wp.vec4(22.0, 21.0, 20.0, 19.0), True)
+
+        v[::-2] = vec2(23.0, 24.0)
+        wp.expect_eq(v == wp.vec4(22.0, 24.0, 20.0, 23.0), True)
+
+        v[1::-2] = vec1(25.0)
+        wp.expect_eq(v == wp.vec4(22.0, 25.0, 20.0, 23.0), True)
+
+        v[1:] += vec3(26.0, 27.0, 28.0)
+        wp.expect_eq(v == wp.vec4(22.0, 51.0, 47.0, 51.0), True)
+
+        v[:-1] -= vec3(29.0, 30.0, 31.0)
+        wp.expect_eq(v == wp.vec4(-7.0, 21.0, 16.0, 51.0), True)
+
+        v[:] %= vec4(32.0, 33.0, 34.0, 35.0)
+        wp.expect_eq(v == wp.vec4(-7.0, 21.0, 16.0, 16.0), True)
+
+    @wp.kernel(module="unique")
+    def kernel():
+        fn()
+
+    wp.launch(kernel, 1, device=device)
+    wp.synchronize()
+    fn()
+
+
+def test_vec_assign_inplace_errors(test, device):
+    @wp.kernel
+    def kernel_1():
+        v = wp.vec4(1.0, 2.0, 3.0, 4.0)
+        v[1:] = wp.vec3d(wp.float64(5.0), wp.float64(6.0), wp.float64(7.0))
+
+    with test.assertRaisesRegex(
+        ValueError,
+        r"The provided vector is expected to be of length 3 with dtype float32.$",
+    ):
+        wp.launch(kernel_1, dim=1, device=device)
+
+    @wp.kernel
+    def kernel_2():
+        v = wp.vec4(1.0, 2.0, 3.0, 4.0)
+        v[1:] = wp.float64(5.0)
+
+    with test.assertRaisesRegex(
+        ValueError,
+        r"The provided value is expected to be a vector of length 3, with dtype float32.$",
+    ):
+        wp.launch(kernel_2, dim=1, device=device)
+
+    @wp.kernel
+    def kernel_3():
+        v = wp.vec4(1.0, 2.0, 3.0, 4.0)
+        v[1:] = wp.mat22(5.0, 6.0, 7.0, 8.0)
+
+    with test.assertRaisesRegex(
+        ValueError,
+        r"The provided value is expected to be a vector of length 3, with dtype float32.$",
+    ):
+        wp.launch(kernel_3, dim=1, device=device)
+
+    @wp.kernel
+    def kernel_4():
+        v = wp.vec4(1.0, 2.0, 3.0, 4.0)
+        v[1:] = wp.vec2(5.0, 6.0)
+
+    with test.assertRaisesRegex(
+        ValueError,
+        r"The length of the provided vector \(2\) isn't compatible with the given slice \(expected 3\).$",
+    ):
+        wp.launch(kernel_4, dim=1, device=device)
+
+
+def test_vec_slicing_assign_backward(test, device):
+    @wp.kernel(module="unique")
+    def kernel(arr_x: wp.array(dtype=wp.vec2), arr_y: wp.array(dtype=wp.vec4)):
+        i = wp.tid()
+
+        y = arr_y[i]
+
+        y[:2] = arr_x[i]
+        y[1:-1] += 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.vec4, 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()
 
 
@@ -1248,7 +1390,6 @@ add_function_test(TestVec, "test_length_mismatch", test_length_mismatch, devices
 add_function_test(TestVec, "test_vector_len", test_vector_len, devices=devices)
 add_function_test(TestVec, "test_vec_extract", test_vec_extract, devices=devices)
 add_function_test(TestVec, "test_vec_assign", test_vec_assign, devices=devices)
-add_function_test(TestVec, "test_vec_assign_copy", test_vec_assign_copy, devices=devices)
 add_function_test(TestVec, "test_vec_array_extract", test_vec_array_extract, devices=devices)
 add_function_test(TestVec, "test_vec_array_assign", test_vec_array_assign, devices=devices)
 add_function_test(TestVec, "test_vec_add_inplace", test_vec_add_inplace, devices=devices)
@@ -1256,6 +1397,10 @@ add_function_test(TestVec, "test_vec_sub_inplace", test_vec_sub_inplace, devices
 add_function_test(TestVec, "test_vec_array_add_inplace", test_vec_array_add_inplace, devices=devices)
 add_function_test(TestVec, "test_vec_array_sub_inplace", test_vec_array_sub_inplace, devices=devices)
 add_function_test(TestVec, "test_scalar_vec_div", test_scalar_vec_div, devices=devices)
+add_function_test(TestVec, "test_vec_indexing_assign", test_vec_indexing_assign, devices=devices)
+add_function_test(TestVec, "test_vec_slicing_assign", test_vec_slicing_assign, devices=devices)
+add_function_test(TestVec, "test_vec_assign_inplace_errors", test_vec_assign_inplace_errors, devices=devices)
+add_function_test(TestVec, "test_vec_slicing_assign_backward", test_vec_slicing_assign_backward, devices=devices)
 
 
 if __name__ == "__main__":

warp/tests/test_vec_assign_copy.py

@@ -0,0 +1,143 @@
+# 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 vec_assign_subscript(x: wp.array(dtype=float), y: wp.array(dtype=wp.vec3)):
+    i = wp.tid()
+
+    a = wp.vec3()
+    a[0] = 1.0 * x[i]
+    a[1] = 2.0 * x[i]
+    a[2] = 3.0 * x[i]
+    y[i] = a
+
+
+@wp.kernel
+def vec_assign_attribute(x: wp.array(dtype=float), y: wp.array(dtype=wp.vec3)):
+    i = wp.tid()
+
+    a = wp.vec3()
+    a.x = 1.0 * x[i]
+    a.y = 2.0 * x[i]
+    a.z = 3.0 * x[i]
+    y[i] = a
+
+
+def test_vec_assign(test, device):
+    def run(kernel):
+        x = wp.ones(1, dtype=float, requires_grad=True, device=device)
+        y = wp.zeros(1, dtype=wp.vec3, 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]], dtype=float))
+        assert_np_equal(x.grad.numpy(), np.array([6.0], dtype=float))
+
+    run(vec_assign_subscript)
+    run(vec_assign_attribute)
+
+
+def test_vec_assign_copy(test, device):
+    @wp.kernel(module="unique")
+    def vec_assign_overwrite(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=wp.vec3)):
+        tid = wp.tid()
+
+        a = wp.vec3()
+        b = x[tid]
+        a = b
+        a[1] = 3.0
+
+        y[tid] = a
+
+    x = wp.ones(1, dtype=wp.vec3, device=device, requires_grad=True)
+    y = wp.zeros(1, dtype=wp.vec3, device=device, requires_grad=True)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(vec_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]], dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array([[1.0, 0.0, 1.0]], dtype=float))
+
+
+def test_vec_slicing_assign_backward(test, device):
+    @wp.kernel(module="unique")
+    def kernel(arr_x: wp.array(dtype=wp.vec2), arr_y: wp.array(dtype=wp.vec4)):
+        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.vec4, 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 TestVecAssignCopy(unittest.TestCase):
+    pass
+
+
+add_function_test(TestVecAssignCopy, "test_vec_assign", test_vec_assign, devices=devices)
+add_function_test(TestVecAssignCopy, "test_vec_assign_copy", test_vec_assign_copy, devices=devices)
+add_function_test(
+    TestVecAssignCopy, "test_vec_slicing_assign_backward", test_vec_slicing_assign_backward, devices=devices
+)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=True)