warp-lang 1.6.2__py3-none-macosx_10_13_universal2.whl → 1.7.1__py3-none-macosx_10_13_universal2.whl

warp/tests/test_mat_scalar_ops.py

@@ -334,19 +334,19 @@ def test_constructors(test, device, dtype, register_kernels=False):
         outcomponents: wp.array(dtype=wptype),
     ):
         # multiply outputs by 2 so we've got something to backpropagate:
-        m2result = wptype(2) * mat22(vec2(input[0], input[2]), vec2(input[1], input[3]))
-        m3result = wptype(2) * mat33(
+        m2result = wptype(2) * wp.matrix_from_cols(vec2(input[0], input[2]), vec2(input[1], input[3]))
+        m3result = wptype(2) * wp.matrix_from_cols(
             vec3(input[4], input[7], input[10]),
             vec3(input[5], input[8], input[11]),
             vec3(input[6], input[9], input[12]),
         )
-        m4result = wptype(2) * mat44(
+        m4result = wptype(2) * wp.matrix_from_cols(
             vec4(input[13], input[17], input[21], input[25]),
             vec4(input[14], input[18], input[22], input[26]),
             vec4(input[15], input[19], input[23], input[27]),
             vec4(input[16], input[20], input[24], input[28]),
         )
-        m5result = wptype(2) * mat55(
+        m5result = wptype(2) * wp.matrix_from_cols(
             vec5(input[29], input[34], input[39], input[44], input[49]),
             vec5(input[30], input[35], input[40], input[45], input[50]),
             vec5(input[31], input[36], input[41], input[46], input[51]),

warp/tests/test_overwrite.py

@@ -23,8 +23,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()  # For wp.context.runtime.core.is_cutlass_enabled()
-
 # kernels are defined in the global scope, to ensure wp.Kernel objects are not GC'ed in the MGPU case
 # kernel args are assigned array modes during codegen, so wp.Kernel objects generated during codegen
 # must be preserved for overwrite tracking to function
@@ -378,62 +376,6 @@ def test_copy(test, device):
         wp.config.verify_autograd_array_access = saved_verify_autograd_array_access_setting
 
 
-# wp.matmul uses wp.record_func. Ensure array modes are propagated correctly.
-def test_matmul(test, device):
-    if device.is_cuda and not wp.context.runtime.core.is_cutlass_enabled():
-        test.skipTest("Warp was not built with CUTLASS support")
-
-    saved_verify_autograd_array_access_setting = wp.config.verify_autograd_array_access
-    try:
-        wp.config.verify_autograd_array_access = True
-
-        a = wp.ones((3, 3), dtype=float, requires_grad=True, device=device)
-        b = wp.ones_like(a)
-        c = wp.ones_like(a)
-        d = wp.zeros_like(a)
-
-        tape = wp.Tape()
-
-        with tape:
-            wp.matmul(a, b, c, d)
-
-        test.assertEqual(a._is_read, True)
-        test.assertEqual(b._is_read, True)
-        test.assertEqual(c._is_read, True)
-        test.assertEqual(d._is_read, False)
-
-    finally:
-        wp.config.verify_autograd_array_access = saved_verify_autograd_array_access_setting
-
-
-# wp.batched_matmul uses wp.record_func. Ensure array modes are propagated correctly.
-def test_batched_matmul(test, device):
-    if device.is_cuda and not wp.context.runtime.core.is_cutlass_enabled():
-        test.skipTest("Warp was not built with CUTLASS support")
-
-    saved_verify_autograd_array_access_setting = wp.config.verify_autograd_array_access
-    try:
-        wp.config.verify_autograd_array_access = True
-
-        a = wp.ones((1, 3, 3), dtype=float, requires_grad=True, device=device)
-        b = wp.ones_like(a)
-        c = wp.ones_like(a)
-        d = wp.zeros_like(a)
-
-        tape = wp.Tape()
-
-        with tape:
-            wp.batched_matmul(a, b, c, d)
-
-        test.assertEqual(a._is_read, True)
-        test.assertEqual(b._is_read, True)
-        test.assertEqual(c._is_read, True)
-        test.assertEqual(d._is_read, False)
-
-    finally:
-        wp.config.verify_autograd_array_access = saved_verify_autograd_array_access_setting
-
-
 # write after read warning with in-place operators within a kernel
 def test_in_place_operators_warning(test, device):
     saved_verify_autograd_array_access_setting = wp.config.verify_autograd_array_access
@@ -593,8 +535,6 @@ add_function_test(TestOverwrite, "test_views", test_views, devices=devices)
 add_function_test(TestOverwrite, "test_reset", test_reset, devices=devices)
 
 add_function_test(TestOverwrite, "test_copy", test_copy, devices=devices)
-add_function_test(TestOverwrite, "test_matmul", test_matmul, devices=devices, check_output=False)
-add_function_test(TestOverwrite, "test_batched_matmul", test_batched_matmul, devices=devices, check_output=False)
 add_function_test(TestOverwrite, "test_atomic_operations", test_atomic_operations, devices=devices)
 
 # Some warning are only issued during codegen, and codegen only runs on cuda_0 in the MGPU case.

warp/tests/test_quat.py

@@ -1205,7 +1205,6 @@ def test_quat_to_matrix(test, device, dtype, register_kernels=False):
 
     wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
     quat = wp.types.quaternion(dtype=wptype)
-    mat3 = wp.types.matrix(shape=(3, 3), dtype=wptype)
     vec3 = wp.types.vector(length=3, dtype=wptype)
 
     def check_quat_to_matrix(
@@ -1239,7 +1238,7 @@ def test_quat_to_matrix(test, device, dtype, register_kernels=False):
                 wptype(1),
             ),
         )
-        result_manual = mat3(xaxis, yaxis, zaxis)
+        result_manual = wp.matrix_from_cols(xaxis, yaxis, zaxis)
 
         idx = 0
         for i in range(3):
@@ -1711,18 +1710,31 @@ def test_quat_rpy_grad(test, device, dtype, register_kernels=False):
 def test_quat_from_matrix(test, device, dtype, register_kernels=False):
     wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
     mat33 = wp.types.matrix((3, 3), wptype)
+    mat44 = wp.types.matrix((4, 4), wptype)
     quat = wp.types.quaternion(wptype)
 
     def quat_from_matrix(m: wp.array2d(dtype=wptype), loss: wp.array(dtype=wptype), idx: int):
         tid = wp.tid()
 
-        matrix = mat33(
-            m[tid, 0], m[tid, 1], m[tid, 2], m[tid, 3], m[tid, 4], m[tid, 5], m[tid, 6], m[tid, 7], m[tid, 8]
+        # fmt: off
+        m3 = mat33(
+            m[tid, 0], m[tid, 1], m[tid, 2],
+            m[tid, 3], m[tid, 4], m[tid, 5],
+            m[tid, 6], m[tid, 7], m[tid, 8],
         )
+        q1 = wp.quat_from_matrix(m3)
 
-        q = wp.quat_from_matrix(matrix)
+        m4 = mat44(
+            m[tid, 0], m[tid, 1], m[tid, 2], wptype(0.0),
+            m[tid, 3], m[tid, 4], m[tid, 5], wptype(0.0),
+            m[tid, 6], m[tid, 7], m[tid, 8], wptype(0.0),
+            wptype(0.0), wptype(0.0), wptype(0.0), wptype(1.0),
+        )
+        q2 = wp.quat_from_matrix(m4)
+        # fmt: on
 
-        wp.atomic_add(loss, 0, q[idx])
+        wp.expect_eq(q1, q2)
+        wp.atomic_add(loss, 0, q1[idx])
 
     def quat_from_matrix_forward(mats: wp.array2d(dtype=wptype), loss: wp.array(dtype=wptype), idx: int):
         tid = wp.tid()
@@ -1891,113 +1903,6 @@ def test_quat_identity(test, device, dtype, register_kernels=False):
     assert_np_equal(output.numpy(), expected)
 
 
-############################################################
-
-
-def test_quat_assign(test, device, dtype, register_kernels=False):
-    np_type = np.dtype(dtype)
-    wp_type = wp.types.np_dtype_to_warp_type[np_type]
-
-    quat = wp.types.quaternion(dtype=wp_type)
-
-    def quattest_read_write_store(x: wp.array(dtype=wp_type), a: wp.array(dtype=quat)):
-        tid = wp.tid()
-
-        t = a[tid]
-        t[0] = x[tid]
-        a[tid] = t
-
-    def quattest_in_register(x: wp.array(dtype=wp_type), a: wp.array(dtype=quat)):
-        tid = wp.tid()
-
-        g = wp_type(0.0)
-        q = a[tid]
-        g = q[0] + wp_type(2.0) * q[1] + wp_type(3.0) * q[2] + wp_type(4.0) * q[3]
-        x[tid] = g
-
-    def quattest_in_register_overwrite(x: wp.array(dtype=quat), a: wp.array(dtype=quat)):
-        tid = wp.tid()
-
-        f = quat()
-        a_quat = a[tid]
-        f = a_quat
-        f[1] = wp_type(3.0)
-
-        x[tid] = f
-
-    def quattest_component(x: wp.array(dtype=quat), y: wp.array(dtype=wp_type)):
-        i = wp.tid()
-
-        a = quat()
-        a.x = wp_type(1.0) * y[i]
-        a.y = wp_type(2.0) * y[i]
-        a.z = wp_type(3.0) * y[i]
-        a.w = wp_type(4.0) * y[i]
-        x[i] = a
-
-    kernel_read_write_store = getkernel(quattest_read_write_store, suffix=dtype.__name__)
-    kernel_in_register = getkernel(quattest_in_register, suffix=dtype.__name__)
-    kernel_in_register_overwrite = getkernel(quattest_in_register_overwrite, suffix=dtype.__name__)
-    kernel_component = getkernel(quattest_component, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    a = wp.ones(1, dtype=quat, device=device, requires_grad=True)
-    x = wp.full(1, value=2.0, dtype=wp_type, device=device, requires_grad=True)
-
-    tape = wp.Tape()
-    with tape:
-        wp.launch(kernel_read_write_store, dim=1, inputs=[x, a], device=device)
-
-    tape.backward(grads={a: wp.ones_like(a, requires_grad=False)})
-
-    assert_np_equal(a.numpy(), np.array([[2.0, 1.0, 1.0, 1.0]], dtype=np_type))
-    assert_np_equal(x.grad.numpy(), np.array([1.0], dtype=np_type))
-
-    tape.reset()
-
-    a = wp.ones(1, dtype=quat, device=device, requires_grad=True)
-    x = wp.zeros(1, dtype=wp_type, device=device, requires_grad=True)
-
-    with tape:
-        wp.launch(kernel_in_register, dim=1, inputs=[x, a], device=device)
-
-    tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
-
-    assert_np_equal(x.numpy(), np.array([10.0], dtype=np_type))
-    assert_np_equal(a.grad.numpy(), np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np_type))
-
-    tape.reset()
-
-    x = wp.zeros(1, dtype=quat, requires_grad=True)
-    y = wp.ones(1, dtype=wp_type, requires_grad=True)
-
-    tape = wp.Tape()
-    with tape:
-        wp.launch(kernel_component, dim=1, inputs=[x, y])
-
-    tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
-
-    assert_np_equal(x.numpy(), np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np_type))
-    assert_np_equal(y.grad.numpy(), np.array([10.0], dtype=np_type))
-
-    x = wp.zeros(1, dtype=quat, device=device, requires_grad=True)
-    a = wp.ones(1, dtype=quat, device=device, requires_grad=True)
-
-    tape = wp.Tape()
-    with tape:
-        wp.launch(kernel_in_register_overwrite, dim=1, inputs=[x, a], device=device)
-
-    tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
-
-    assert_np_equal(x.numpy(), np.array([[1.0, 3.0, 1.0, 1.0]], dtype=np_type))
-    assert_np_equal(a.grad.numpy(), np.array([[1.0, 0.0, 1.0, 1.0]], dtype=np_type))
-
-
-############################################################
-
-
 def test_quat_euler_conversion(test, device, dtype, register_kernels=False):
     rng = np.random.default_rng(123)
     N = 3
@@ -2077,6 +1982,12 @@ def test_constructor_default():
     wp.expect_eq(qeye[2], 0.0)
     wp.expect_eq(qeye[3], 1.0)
 
+    qlit = wp.quaternion(1.0, 2.0, 3.0, 4.0, dtype=float)
+    wp.expect_eq(qlit[0], 1.0)
+    wp.expect_eq(qlit[1], 2.0)
+    wp.expect_eq(qlit[2], 3.0)
+    wp.expect_eq(qlit[3], 4.0)
+
 
 def test_py_arithmetic_ops(test, device, dtype):
     wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
@@ -2128,56 +2039,348 @@ def test_quat_len(test, device):
 
 
 @wp.kernel
-def vector_augassign_kernel(
-    a: wp.array(dtype=wp.quat), b: wp.array(dtype=wp.quat), c: wp.array(dtype=wp.quat), d: wp.array(dtype=wp.quat)
-):
+def quat_extract_subscript(x: wp.array(dtype=wp.quat), 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]
+    y[tid] = b
+
+
+""" TODO: rhs attribute indexing
+@wp.kernel
+def quat_extract_attribute(x: wp.array(dtype=wp.quat), y: wp.array(dtype=float)):
+    tid = wp.tid()
+
+    a = x[tid]
+    b = a.x + float(2.0) * a.y + 3.0 * a.z + 4.0 * a.w
+    y[tid] = b
+"""
+
+
+def test_quat_extract(test, device):
+    def run(kernel):
+        x = wp.ones(1, dtype=wp.quat, requires_grad=True, device=device)
+        y = wp.zeros(1, dtype=float, 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([10.0], dtype=float))
+        assert_np_equal(x.grad.numpy(), np.array([[1.0, 2.0, 3.0, 4.0]], dtype=float))
+
+    run(quat_extract_subscript)
+    # run(quat_extract_attribute)
+
+
+@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()
 
-    q1 = wp.quat()
-    q2 = b[i]
+    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
+
 
-    q1[0] += q2[0]
-    q1[1] += q2[1]
-    q1[2] += q2[2]
-    q1[3] += q2[3]
+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)
 
-    a[i] = q1
+        tape = wp.Tape()
+        with tape:
+            wp.launch(kernel, 1, inputs=[x], outputs=[y], device=device)
 
-    q3 = wp.quat()
-    q4 = d[i]
+        y.grad = wp.ones_like(y)
+        tape.backward()
 
-    q3[0] += q4[0]
-    q3[1] += q4[1]
-    q3[2] += q4[2]
-    q3[3] += q4[3]
+        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))
 
-    c[i] = q1
+    run(quat_assign_subscript)
+    run(quat_assign_attribute)
 
 
-def test_vector_augassign(test, device):
-    N = 3
+def test_quat_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 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))
+
+    finally:
+        wp.config.enable_vector_component_overwrites = saved_enable_vector_component_overwrites_setting
+
 
-    a = wp.zeros(N, dtype=wp.quat, requires_grad=True)
-    b = wp.ones(N, dtype=wp.quat, requires_grad=True)
+@wp.kernel
+def quat_array_extract_subscript(x: wp.array2d(dtype=wp.quat), 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]
+    y[i, j] = 1.0 * a + 2.0 * b + 3.0 * c + 4.0 * d
+
+
+""" TODO: rhs attribute indexing
+@wp.kernel
+def quat_array_extract_attribute(x: wp.array2d(dtype=wp.quat), y: wp.array2d(dtype=float)):
+    i, j = wp.tid()
+    a = x[i, j].x
+    b = x[i, j].y
+    c = x[i, j].z
+    d = x[i, j].w
+    y[i, j] = 1.0 * a + 2.0 * b + 3.0 * c + 4.0 * d
+"""
+
+
+def test_quat_array_extract(test, device):
+    def run(kernel):
+        x = wp.ones((1, 1), dtype=wp.quat, requires_grad=True, device=device)
+        y = wp.zeros((1, 1), dtype=float, requires_grad=True, device=device)
+
+        tape = wp.Tape()
+        with tape:
+            wp.launch(kernel, (1, 1), inputs=[x], outputs=[y], device=device)
+
+        y.grad = wp.ones_like(y)
+        tape.backward()
+
+        assert_np_equal(y.numpy(), np.array([[10.0]], dtype=float))
+        assert_np_equal(x.grad.numpy(), np.array([[[1.0, 2.0, 3.0, 4.0]]], dtype=float))
+
+    run(quat_array_extract_subscript)
+    # run(quat_array_extract_attribute)
+
+
+@wp.kernel
+def quat_array_assign_subscript(x: wp.array2d(dtype=float), y: wp.array2d(dtype=wp.quat)):
+    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]
+
+
+@wp.kernel
+def quat_array_assign_attribute(x: wp.array2d(dtype=float), y: wp.array2d(dtype=wp.quat)):
+    i, j = wp.tid()
+
+    y[i, j].x = 1.0 * x[i, j]
+    y[i, j].y = 2.0 * x[i, j]
+    y[i, j].z = 3.0 * x[i, j]
+    y[i, j].w = 4.0 * x[i, j]
+
+
+def test_quat_array_assign(test, device):
+    def run(kernel):
+        x = wp.ones((1, 1), dtype=float, requires_grad=True, device=device)
+        y = wp.zeros((1, 1), dtype=wp.quat, requires_grad=True, device=device)
+
+        tape = wp.Tape()
+        with tape:
+            wp.launch(kernel, (1, 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))
+        # TODO: gradient propagation for in-place array assignment
+        # assert_np_equal(x.grad.numpy(), np.array([[10.0]], dtype=float))
+
+    run(quat_array_assign_subscript)
+    run(quat_array_assign_attribute)
+
+
+@wp.kernel
+def quat_add_inplace_subscript(x: wp.array(dtype=wp.quat), y: wp.array(dtype=wp.quat)):
+    i = wp.tid()
+
+    a = wp.quat()
+    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]
+
+    y[i] = a
+
+
+""" TODO: rhs attribute indexing
+@wp.kernel
+def quat_add_inplace_attribute(x: wp.array(dtype=wp.quat), y: wp.array(dtype=wp.quat)):
+    i = wp.tid()
+
+    a = wp.quat()
+    b = x[i]
+
+    a.x += 1.0 * b.x
+    a.y += 2.0 * b.y
+    a.z += 3.0 * b.z
+    a.w += 4.0 * b.w
+
+    y[i] = a
+"""
+
+
+def test_quat_add_inplace(test, device):
+    def run(kernel):
+        x = wp.ones(1, dtype=wp.quat, 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)
 
-    c = wp.zeros(N, dtype=wp.quat, requires_grad=True)
-    d = wp.ones(N, dtype=wp.quat, requires_grad=True)
+        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([[1.0, 2.0, 3.0, 4.0]], dtype=float))
+
+    run(quat_add_inplace_subscript)
+    # run(quat_add_inplace_attribute)
+
+
+@wp.kernel
+def quat_sub_inplace_subscript(x: wp.array(dtype=wp.quat), y: wp.array(dtype=wp.quat)):
+    i = wp.tid()
+
+    a = wp.quat()
+    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]
+
+    y[i] = a
+
+
+""" TODO: rhs attribute indexing
+@wp.kernel
+def quat_sub_inplace_attribute(x: wp.array(dtype=wp.quat), y: wp.array(dtype=wp.quat)):
+    i = wp.tid()
+
+    a = wp.quat()
+    b = x[i]
+
+    a.x -= 1.0 * b.x
+    a.y -= 2.0 * b.y
+    a.z -= 3.0 * b.z
+    a.w -= 4.0 * b.w
+
+    y[i] = a
+"""
+
+
+def test_quat_sub_inplace(test, device):
+    def run(kernel):
+        x = wp.ones(1, dtype=wp.quat, 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([[-1.0, -2.0, -3.0, -4.0]], dtype=float))
+
+    run(quat_sub_inplace_subscript)
+    # run(quat_sub_inplace_attribute)
+
+
+@wp.kernel
+def quat_array_add_inplace(x: wp.array(dtype=wp.quat), y: wp.array(dtype=wp.quat)):
+    i = wp.tid()
+
+    y[i] += x[i]
+
+
+def test_quat_array_add_inplace(test, device):
+    x = wp.ones(1, dtype=wp.quat, requires_grad=True, device=device)
+    y = wp.zeros(1, dtype=wp.quat, requires_grad=True, device=device)
 
     tape = wp.Tape()
     with tape:
-        wp.launch(vector_augassign_kernel, N, inputs=[a, b, c, d])
+        wp.launch(quat_array_add_inplace, 1, inputs=[x], outputs=[y], device=device)
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
 
-    tape.backward(grads={a: wp.ones_like(a), c: wp.ones_like(c)})
+    assert_np_equal(y.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]], dtype=float))
 
-    assert_np_equal(a.numpy(), wp.ones_like(a).numpy())
-    assert_np_equal(a.grad.numpy(), wp.ones_like(a).numpy())
-    assert_np_equal(b.grad.numpy(), wp.ones_like(a).numpy())
 
-    assert_np_equal(c.numpy(), -wp.ones_like(c).numpy())
-    assert_np_equal(c.grad.numpy(), wp.ones_like(c).numpy())
-    assert_np_equal(d.grad.numpy(), -wp.ones_like(d).numpy())
+""" TODO: quat negation operator
+@wp.kernel
+def quat_array_sub_inplace(x: wp.array(dtype=wp.quat), y: wp.array(dtype=wp.quat)):
+    i = wp.tid()
+
+    y[i] -= x[i]
 
 
+def test_quat_array_sub_inplace(test, device):
+    x = wp.ones(1, dtype=wp.quat, requires_grad=True, device=device)
+    y = wp.zeros(1, dtype=wp.quat, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(quat_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]], dtype=float))
+    assert_np_equal(x.grad.numpy(), np.array([[-1.0, -1.0, -1.0, -1.0]], dtype=float))
+"""
+
 devices = get_test_devices()
 
 
@@ -2275,18 +2478,20 @@ for dtype in np_float_types:
         devices=devices,
         dtype=dtype,
     )
-    add_function_test_register_kernel(
-        TestQuat,
-        f"test_quat_assign_{dtype.__name__}",
-        test_quat_assign,
-        devices=devices,
-        dtype=dtype,
-    )
     add_function_test(
         TestQuat, f"test_py_arithmetic_ops_{dtype.__name__}", test_py_arithmetic_ops, devices=None, dtype=dtype
     )
 
 add_function_test(TestQuat, "test_quat_len", test_quat_len, devices=devices)
+add_function_test(TestQuat, "test_quat_extract", test_quat_extract, devices=devices)
+add_function_test(TestQuat, "test_quat_assign", test_quat_assign, devices=devices)
+add_function_test(TestQuat, "test_quat_assign_copy", test_quat_assign_copy, devices=devices)
+add_function_test(TestQuat, "test_quat_array_extract", test_quat_array_extract, devices=devices)
+add_function_test(TestQuat, "test_quat_array_assign", test_quat_array_assign, devices=devices)
+add_function_test(TestQuat, "test_quat_add_inplace", test_quat_add_inplace, devices=devices)
+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)
 
 
 if __name__ == "__main__":