warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/tests/test_mat.py

@@ -5,9 +5,12 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import unittest
+
 import numpy as np
+
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -19,37 +22,24 @@ np_signed_int_types = [
     np.byte,
 ]
 
-np_unsigned_int_types = [
-    np.uint8,
-    np.uint16,
-    np.uint32,
-    np.uint64,
-    np.ubyte,
-]
-
-np_int_types = np_signed_int_types + np_unsigned_int_types
-
 np_float_types = [np.float16, np.float32, np.float64]
 
-np_scalar_types = np_int_types + np_float_types
-
 
-def randvals(shape, dtype):
+def randvals(rng, shape, dtype):
     if dtype in np_float_types:
-        return np.random.randn(*shape).astype(dtype)
+        return rng.standard_normal(size=shape).astype(dtype)
     elif dtype in [np.int8, np.uint8, np.byte, np.ubyte]:
-        return np.random.randint(1, 3, size=shape, dtype=dtype)
-    return np.random.randint(1, 5, size=shape, dtype=dtype)
+        return rng.integers(1, high=3, size=shape, dtype=dtype)
+    return rng.integers(1, high=5, size=shape, dtype=dtype)
 
 
 kernel_cache = dict()
 
 
 def getkernel(func, suffix=""):
-    module = wp.get_module(func.__module__)
     key = func.__name__ + "_" + suffix
     if key not in kernel_cache:
-        kernel_cache[key] = wp.Kernel(func=func, key=key, module=module)
+        kernel_cache[key] = wp.Kernel(func=func, key=key)
     return kernel_cache[key]
 
 
@@ -63,324 +53,224 @@ def get_select_kernel(dtype):
 
     return getkernel(output_select_kernel_fn, suffix=dtype.__name__)
 
+    wp.launch(kernel, dim=1, inputs=[])
 
-def test_arrays(test, device, dtype):
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-    mat32 = wp.types.matrix(shape=(3, 2), dtype=wptype)
 
-    np.random.seed(123)
+def test_anon_constructor_error_shape_keyword_missing(test, device):
+    @wp.kernel
+    def kernel():
+        wp.matrix(1.0, 2.0, 3.0)
 
-    v2_np = randvals([10, 2, 2], dtype)
-    v3_np = randvals([10, 3, 3], dtype)
-    v4_np = randvals([10, 4, 4], dtype)
-    v5_np = randvals([10, 5, 5], dtype)
-    v32_np = randvals([10, 3, 2], dtype)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"shape keyword must be specified when calling matrix\(\) function$",
+    ):
+        wp.launch(
+            kernel,
+            dim=1,
+            inputs=[],
+            device=device,
+        )
 
-    v2 = wp.array(v2_np, dtype=mat22, requires_grad=True, device=device)
-    v3 = wp.array(v3_np, dtype=mat33, requires_grad=True, device=device)
-    v4 = wp.array(v4_np, dtype=mat44, requires_grad=True, device=device)
-    v5 = wp.array(v5_np, dtype=mat55, requires_grad=True, device=device)
-    v32 = wp.array(v32_np, dtype=mat32, requires_grad=True, device=device)
 
-    assert_np_equal(v2.numpy(), v2_np, tol=1.0e-6)
-    assert_np_equal(v3.numpy(), v3_np, tol=1.0e-6)
-    assert_np_equal(v4.numpy(), v4_np, tol=1.0e-6)
-    assert_np_equal(v5.numpy(), v5_np, tol=1.0e-6)
-    assert_np_equal(v32.numpy(), v32_np, tol=1.0e-6)
+def test_anon_constructor_error_dtype_keyword_missing(test, device):
+    @wp.kernel
+    def kernel():
+        wp.matrix(shape=(3, 3))
 
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"matrix\(\) must have dtype as a keyword argument if it has no " r"positional arguments$",
+    ):
+        wp.launch(
+            kernel,
+            dim=1,
+            inputs=[],
+            device=device,
+        )
 
-    v2 = wp.array(v2_np, dtype=mat22, requires_grad=True, device=device)
-    v3 = wp.array(v3_np, dtype=mat33, requires_grad=True, device=device)
-    v4 = wp.array(v4_np, dtype=mat44, requires_grad=True, device=device)
 
-    assert_np_equal(v2.numpy(), v2_np, tol=1.0e-6)
-    assert_np_equal(v3.numpy(), v3_np, tol=1.0e-6)
-    assert_np_equal(v4.numpy(), v4_np, tol=1.0e-6)
+def test_anon_constructor_error_shape_mismatch(test, device):
+    @wp.kernel
+    def kernel():
+        wp.matrix(
+            wp.matrix(shape=(1, 2), dtype=float),
+            shape=(3, 4),
+            dtype=float,
+        )
 
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"Incompatible matrix sizes for casting copy constructor, " r"\(3, 4\) vs \(1, 2\)$",
+    ):
+        wp.launch(
+            kernel,
+            dim=1,
+            inputs=[],
+            device=device,
+        )
 
-def test_constants(test, device, dtype, register_kernels=False):
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-    mat32 = wp.types.matrix(shape=(3, 2), dtype=wptype)
 
-    cm22 = wp.constant(mat22(22))
-    cm33 = wp.constant(mat33(33))
-    cm44 = wp.constant(mat44(44))
-    cm55 = wp.constant(mat55(55))
-    cm32 = wp.constant(mat32(32))
+def test_anon_constructor_error_invalid_arg_count(test, device):
+    @wp.kernel
+    def kernel():
+        wp.matrix(1.0, 2.0, 3.0, shape=(2, 2), dtype=float)
 
-    def check_matrix_constants():
-        wp.expect_eq(cm22, mat22(wptype(22)))
-        wp.expect_eq(cm33, mat33(wptype(33)))
-        wp.expect_eq(cm44, mat44(wptype(44)))
-        wp.expect_eq(cm55, mat55(wptype(55)))
-        wp.expect_eq(cm32, mat32(wptype(32)))
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"Wrong number of arguments for matrix\(\) function, must initialize "
+        r"with either a scalar value, or m\*n values$",
+    ):
+        wp.launch(
+            kernel,
+            dim=1,
+            inputs=[],
+            device=device,
+        )
 
-    kernel = getkernel(check_matrix_constants, suffix=dtype.__name__)
 
-    if register_kernels:
-        return
+def test_tpl_constructor_error_incompatible_sizes(test, device):
+    @wp.kernel
+    def kernel():
+        wp.mat33(wp.mat22(1.0, 2.0, 3.0, 4.0))
 
-    wp.launch(kernel, dim=1, inputs=[])
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"Incompatible matrix sizes for casting copy constructor, " r"\(3, 3\) vs \(2, 2\)$",
+    ):
+        wp.launch(
+            kernel,
+            dim=1,
+            inputs=[],
+            device=device,
+        )
 
 
-def test_constructors(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+def test_tpl_constructor_error_invalid_scalar_type(test, device):
+    @wp.kernel
+    def kernel():
+        wp.mat22(1, 2, 3, 4)
 
-    tol = {
-        np.float16: 1.0e-3,
-        np.float32: 1.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"Wrong scalar type for mat 2,2,<class 'warp.types.float32'> constructor$",
+    ):
+        wp.launch(
+            kernel,
+            dim=1,
+            inputs=[],
+            device=device,
+        )
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    vec2 = wp.types.vector(length=2, dtype=wptype)
-    vec3 = wp.types.vector(length=3, dtype=wptype)
-    vec4 = wp.types.vector(length=4, dtype=wptype)
-    vec5 = wp.types.vector(length=5, dtype=wptype)
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
 
-    output_select_kernel = get_select_kernel(wptype)
+def test_tpl_constructor_error_invalid_vector_count(test, device):
+    @wp.kernel
+    def kernel():
+        wp.mat22(wp.vec3(1.0, 2.0, 3.0))
 
-    def check_scalar_mat_constructor(
-        input: wp.array(dtype=wptype),
-        outcomponents: wp.array(dtype=wptype),
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"Wrong number of vectors when attempting to construct a matrix " r"with column vectors$",
     ):
-        # multiply outputs by 2 so we've got something to backpropagate:
-        m2result = wptype(2) * mat22(input[0])
-        m3result = wptype(2) * mat33(input[0])
-        m4result = wptype(2) * mat44(input[0])
-        m5result = wptype(2) * mat55(input[0])
-
-        idx = 0
-        for i in range(2):
-            for j in range(2):
-                outcomponents[idx] = m2result[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(3):
-                outcomponents[idx] = m3result[i, j]
-                idx = idx + 1
+        wp.launch(
+            kernel,
+            dim=1,
+            inputs=[],
+            device=device,
+        )
 
-        for i in range(4):
-            for j in range(4):
-                outcomponents[idx] = m4result[i, j]
-                idx = idx + 1
 
-        for i in range(5):
-            for j in range(5):
-                outcomponents[idx] = m5result[i, j]
-                idx = idx + 1
+def test_tpl_constructor_error_invalid_vector_shape(test, device):
+    @wp.kernel
+    def kernel():
+        wp.mat22(wp.vec3(1.0, 2.0, 3.0), wp.vec3(4.0, 5.0, 6.0))
 
-    def check_component_mat_constructor(
-        input: wp.array(dtype=wptype),
-        outcomponents: wp.array(dtype=wptype),
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"Wrong vector row count when attempting to construct a matrix " r"with column vectors$",
     ):
-        # multiply outputs by 2 so we've got something to backpropagate:
-        m2result = wptype(2) * mat22(input[0], input[1], input[2], input[3])
-        m3result = wptype(2) * mat33(
-            input[4],
-            input[5],
-            input[6],
-            input[7],
-            input[8],
-            input[9],
-            input[10],
-            input[11],
-            input[12],
-        )
-        m4result = wptype(2) * mat44(
-            input[13],
-            input[14],
-            input[15],
-            input[16],
-            input[17],
-            input[18],
-            input[19],
-            input[20],
-            input[21],
-            input[22],
-            input[23],
-            input[24],
-            input[25],
-            input[26],
-            input[27],
-            input[28],
-        )
-        m5result = wptype(2) * mat55(
-            input[29],
-            input[30],
-            input[31],
-            input[32],
-            input[33],
-            input[34],
-            input[35],
-            input[36],
-            input[37],
-            input[38],
-            input[39],
-            input[40],
-            input[41],
-            input[42],
-            input[43],
-            input[44],
-            input[45],
-            input[46],
-            input[47],
-            input[48],
-            input[49],
-            input[50],
-            input[51],
-            input[52],
-            input[53],
+        wp.launch(
+            kernel,
+            dim=1,
+            inputs=[],
+            device=device,
         )
 
-        idx = 0
-        for i in range(2):
-            for j in range(2):
-                outcomponents[idx] = m2result[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(3):
-                outcomponents[idx] = m3result[i, j]
-                idx = idx + 1
-
-        for i in range(4):
-            for j in range(4):
-                outcomponents[idx] = m4result[i, j]
-                idx = idx + 1
 
-        for i in range(5):
-            for j in range(5):
-                outcomponents[idx] = m5result[i, j]
-                idx = idx + 1
+def test_tpl_constructor_error_invalid_arg_count(test, device):
+    @wp.kernel
+    def kernel():
+        wp.mat22(1.0, 2.0, 3.0)
 
-    def check_vector_mat_constructor(
-        input: wp.array(dtype=wptype),
-        outcomponents: wp.array(dtype=wptype),
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"Wrong number of scalars when attempting to construct a matrix " r"from a list of components$",
     ):
-        # 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(
-            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(
-            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(
-            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]),
-            vec5(input[32], input[37], input[42], input[47], input[52]),
-            vec5(input[33], input[38], input[43], input[48], input[53]),
+        wp.launch(
+            kernel,
+            dim=1,
+            inputs=[],
+            device=device,
         )
 
-        idx = 0
-        for i in range(2):
-            for j in range(2):
-                outcomponents[idx] = m2result[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(3):
-                outcomponents[idx] = m3result[i, j]
-                idx = idx + 1
-
-        for i in range(4):
-            for j in range(4):
-                outcomponents[idx] = m4result[i, j]
-                idx = idx + 1
 
-        for i in range(5):
-            for j in range(5):
-                outcomponents[idx] = m5result[i, j]
-                idx = idx + 1
+def test_tpl_ops_with_anon(test, device):
+    mat22f = wp.mat((2, 2), dtype=float)
 
-    kernel = getkernel(check_scalar_mat_constructor, suffix=dtype.__name__)
-    compkernel = getkernel(check_component_mat_constructor, suffix=dtype.__name__)
-    veckernel = getkernel(check_vector_mat_constructor, suffix=dtype.__name__)
+    m = wp.mat22f(1.0, 2.0, 3.0, 4.0)
+    m += mat22f(2.0, 3.0, 4.0, 5.0)
+    m -= mat22f(3.0, 4.0, 5.0, 6.0)
+    test.assertSequenceEqual(m, ((0.0, 1.0), (2.0, 3.0)))
 
-    if register_kernels:
-        return
+    m = mat22f(1.0, 2.0, 3.0, 4.0)
+    m += wp.mat22f(2.0, 3.0, 4.0, 5.0)
+    m -= wp.mat22f(3.0, 4.0, 5.0, 6.0)
+    test.assertSequenceEqual(m, ((0.0, 1.0), (2.0, 3.0)))
 
-    input = wp.array(randvals([1], dtype), requires_grad=True, device=device)
-    val = input.numpy()[0]
-    outcomponents = wp.zeros(2 * 2 + 3 * 3 + 4 * 4 + 5 * 5, dtype=wptype, requires_grad=True, device=device)
-    out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
 
-    wp.launch(kernel, dim=1, inputs=[input], outputs=[outcomponents], device=device)
+def test_py_arithmetic_ops(test, device, dtype):
+    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
 
-    assert_np_equal(outcomponents.numpy()[:4], 2 * val * np.ones(2 * 2), tol=tol)
-    assert_np_equal(outcomponents.numpy()[4:13], 2 * val * np.ones(3 * 3), tol=tol)
-    assert_np_equal(outcomponents.numpy()[13:29], 2 * val * np.ones(4 * 4), tol=tol)
-    assert_np_equal(outcomponents.numpy()[29:54], 2 * val * np.ones(5 * 5), tol=tol)
+    def make_mat(*args):
+        if wptype in wp.types.int_types:
+            # Cast to the correct integer type to simulate wrapping.
+            return tuple(tuple(wptype._type_(x).value for x in row) for row in args)
 
-    if dtype in np_float_types:
-        for idx in range(len(outcomponents)):
-            tape = wp.Tape()
-            with tape:
-                wp.launch(kernel, dim=1, inputs=[input], outputs=[outcomponents], device=device)
-                wp.launch(output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device)
-            tape.backward(loss=out)
-            test.assertEqual(tape.gradients[input].numpy()[0], 2)
-            tape.zero()
+        return args
 
-    input = wp.array(randvals([2 * 2 + 3 * 3 + 4 * 4 + 5 * 5], dtype), requires_grad=True, device=device)
+    def make_vec(*args):
+        if wptype in wp.types.int_types:
+            # Cast to the correct integer type to simulate wrapping.
+            return tuple(wptype._type_(x).value for x in args)
 
-    wp.launch(compkernel, dim=1, inputs=[input], outputs=[outcomponents], device=device)
-    assert_np_equal(2 * input.numpy(), outcomponents.numpy(), tol=10 * tol)
+        return args
 
-    if dtype in np_float_types:
-        for idx in range(len(outcomponents)):
-            tape = wp.Tape()
-            with tape:
-                wp.launch(compkernel, dim=1, inputs=[input], outputs=[outcomponents], device=device)
-                wp.launch(output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device)
-            tape.backward(loss=out)
-            expectedgrads = np.zeros(len(input))
-            expectedgrads[idx] = 2
-            assert_np_equal(tape.gradients[input].numpy(), expectedgrads)
-            tape.zero()
+    mat_cls = wp.mat((3, 3), wptype)
+    vec_cls = wp.vec(3, wptype)
 
-    wp.launch(veckernel, dim=1, inputs=[input], outputs=[outcomponents], device=device)
-    assert_np_equal(2 * input.numpy(), outcomponents.numpy(), tol=10 * tol)
+    m = mat_cls(((-1, 2, 3), (4, -5, 6), (7, 8, -9)))
+    test.assertSequenceEqual(+m, make_mat((-1, 2, 3), (4, -5, 6), (7, 8, -9)))
+    test.assertSequenceEqual(-m, make_mat((1, -2, -3), (-4, 5, -6), (-7, -8, 9)))
+    test.assertSequenceEqual(m + mat_cls((5, 5, 5) * 3), make_mat((4, 7, 8), (9, 0, 11), (12, 13, -4)))
+    test.assertSequenceEqual(m - mat_cls((5, 5, 5) * 3), make_mat((-6, -3, -2), (-1, -10, 1), (2, 3, -14)))
+    test.assertSequenceEqual(m * vec_cls(5, 5, 5), make_vec(20, 25, 30))
+    test.assertSequenceEqual(m @ vec_cls(5, 5, 5), make_vec(20, 25, 30))
+    test.assertSequenceEqual(vec_cls(5, 5, 5) * m, make_vec(50, 25, 0))
+    test.assertSequenceEqual(vec_cls(5, 5, 5) @ m, make_vec(50, 25, 0))
 
-    if dtype in np_float_types:
-        for idx in range(len(outcomponents)):
-            tape = wp.Tape()
-            with tape:
-                wp.launch(veckernel, dim=1, inputs=[input], outputs=[outcomponents], device=device)
-                wp.launch(output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device)
-            tape.backward(loss=out)
-            expectedgrads = np.zeros(len(input))
-            expectedgrads[idx] = 2
-            assert_np_equal(tape.gradients[input].numpy(), expectedgrads)
-            tape.zero()
+    m = mat_cls(((2, 4, 6), (8, 10, 12), (14, 16, 18)))
+    test.assertSequenceEqual(m * wptype(2), make_mat((4, 8, 12), (16, 20, 24), (28, 32, 36)))
+    test.assertSequenceEqual(wptype(2) * m, make_mat((4, 8, 12), (16, 20, 24), (28, 32, 36)))
+    test.assertSequenceEqual(m / wptype(2), make_mat((1, 2, 3), (4, 5, 6), (7, 8, 9)))
+    test.assertSequenceEqual(wptype(5040) / m, make_mat((2520, 1260, 840), (630, 504, 420), (360, 315, 280)))
+    test.assertSequenceEqual(m * vec_cls(5, 5, 5), make_vec(60, 150, 240))
+    test.assertSequenceEqual(m @ vec_cls(5, 5, 5), make_vec(60, 150, 240))
+    test.assertSequenceEqual(vec_cls(5, 5, 5) * m, make_vec(120, 150, 180))
+    test.assertSequenceEqual(vec_cls(5, 5, 5) @ m, make_vec(120, 150, 180))
 
 
 def test_quat_constructor(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 1.0e-3,
@@ -429,15 +319,15 @@ def test_quat_constructor(test, device, dtype, register_kernels=False):
         return
 
     # translation:
-    p = wp.array(np.random.randn(1, 3).astype(dtype), dtype=vec3, requires_grad=True, device=device)
+    p = wp.array(rng.standard_normal(size=(1, 3)).astype(dtype), dtype=vec3, requires_grad=True, device=device)
 
     # generate a normalized quaternion for the rotation:
-    r = np.random.randn(1, 4)
+    r = rng.standard_normal(size=(1, 4))
     r /= np.linalg.norm(r)
     r = wp.array(r.astype(dtype), dtype=quat, requires_grad=True, device=device)
 
     # scale:
-    s = wp.array(np.random.randn(1, 3).astype(dtype), dtype=vec3, requires_grad=True, device=device)
+    s = wp.array(rng.standard_normal(size=(1, 3)).astype(dtype), dtype=vec3, requires_grad=True, device=device)
 
     # just going to generate the matrix using the constructor, then
     # more manually, and make sure the values/gradients are the same:
@@ -478,95 +368,11 @@ def test_quat_constructor(test, device, dtype, register_kernels=False):
             idx = idx + 1
 
 
-def test_indexing(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 1.0e-3,
-        np.float32: 1.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-
-    output_select_kernel = get_select_kernel(wptype)
-
-    def check_mat_indexing(
-        m2: wp.array(dtype=mat22),
-        m3: wp.array(dtype=mat33),
-        m4: wp.array(dtype=mat44),
-        m5: wp.array(dtype=mat55),
-        outcomponents: wp.array(dtype=wptype),
-    ):
-        # multiply outputs by 2 so we've got something to backpropagate:
-        idx = 0
-        for i in range(2):
-            for j in range(2):
-                outcomponents[idx] = wptype(2) * m2[0][i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(3):
-                outcomponents[idx] = wptype(2) * m3[0][i, j]
-                idx = idx + 1
-
-        for i in range(4):
-            for j in range(4):
-                outcomponents[idx] = wptype(2) * m4[0][i, j]
-                idx = idx + 1
-
-        for i in range(5):
-            for j in range(5):
-                outcomponents[idx] = wptype(2) * m5[0][i, j]
-                idx = idx + 1
-
-    kernel = getkernel(check_mat_indexing, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    m2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    m3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    m4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    m5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    outcomponents = wp.zeros(2 * 2 + 3 * 3 + 4 * 4 + 5 * 5, dtype=wptype, requires_grad=True, device=device)
-
-    wp.launch(kernel, dim=1, inputs=[m2, m3, m4, m5], outputs=[outcomponents], device=device)
-
-    assert_np_equal(outcomponents.numpy()[:4], 2 * m2.numpy().reshape(-1), tol=tol)
-    assert_np_equal(outcomponents.numpy()[4:13], 2 * m3.numpy().reshape(-1), tol=tol)
-    assert_np_equal(outcomponents.numpy()[13:29], 2 * m4.numpy().reshape(-1), tol=tol)
-    assert_np_equal(outcomponents.numpy()[29:54], 2 * m5.numpy().reshape(-1), tol=tol)
-
-    if dtype in np_float_types:
-        idx = 0
-        out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-        for dim, input in [(2, m2), (3, m3), (4, m4), (5, m5)]:
-            for i in range(dim):
-                for j in range(dim):
-                    tape = wp.Tape()
-                    with tape:
-                        wp.launch(kernel, dim=1, inputs=[m2, m3, m4, m5], outputs=[outcomponents], device=device)
-                        wp.launch(
-                            output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device
-                        )
-                    tape.backward(loss=out)
-                    expectedresult = np.zeros((dim, dim), dtype=dtype)
-                    expectedresult[i, j] = 2
-                    assert_np_equal(tape.gradients[input].numpy()[0], expectedresult)
-                    tape.zero()
-                    idx = idx + 1
-
-
-def test_equality(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+def test_negation(test, device, dtype, register_kernels=False):
+    rng = np.random.default_rng(123)
 
     tol = {
-        np.float16: 1.0e-3,
+        np.float16: 1.0e-2,
         np.float32: 1.0e-6,
         np.float64: 1.0e-8,
     }.get(dtype, 0)
@@ -574,1614 +380,85 @@ def test_equality(test, device, dtype, register_kernels=False):
     wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
     mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
     mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-
-    def check_mat_equality():
-        wp.expect_eq(
-            mat22(wptype(1.0), wptype(2.0), wptype(3.0), wptype(4.0)),
-            mat22(wptype(1.0), wptype(2.0), wptype(3.0), wptype(4.0)),
-        )
-        wp.expect_neq(
-            mat22(wptype(1.0), wptype(2.0), wptype(3.0), -wptype(4.0)),
-            mat22(wptype(1.0), wptype(2.0), wptype(3.0), wptype(4.0)),
-        )
-
-        wp.expect_eq(
-            mat33(
-                wptype(1.0),
-                wptype(2.0),
-                wptype(3.0),
-                wptype(4.0),
-                wptype(5.0),
-                wptype(6.0),
-                wptype(7.0),
-                wptype(8.0),
-                wptype(9.0),
-            ),
-            mat33(
-                wptype(1.0),
-                wptype(2.0),
-                wptype(3.0),
-                wptype(4.0),
-                wptype(5.0),
-                wptype(6.0),
-                wptype(7.0),
-                wptype(8.0),
-                wptype(9.0),
-            ),
-        )
-        wp.expect_neq(
-            mat33(
-                wptype(1.0),
-                wptype(2.0),
-                wptype(3.0),
-                wptype(4.0),
-                wptype(5.0),
-                wptype(6.0),
-                wptype(7.0),
-                wptype(8.0),
-                wptype(9.0),
-            ),
-            mat33(
-                wptype(1.0),
-                wptype(2.0),
-                wptype(3.0),
-                -wptype(4.0),
-                wptype(5.0),
-                wptype(6.0),
-                wptype(7.0),
-                wptype(8.0),
-                wptype(9.0),
-            ),
-        )
-
-        wp.expect_eq(
-            mat44(
-                wptype(1.0),
-                wptype(2.0),
-                wptype(3.0),
-                wptype(4.0),
-                wptype(5.0),
-                wptype(6.0),
-                wptype(7.0),
-                wptype(8.0),
-                wptype(9.0),
-                wptype(10.0),
-                wptype(11.0),
-                wptype(12.0),
-                wptype(13.0),
-                wptype(14.0),
-                wptype(15.0),
-                wptype(16.0),
-            ),
-            mat44(
-                wptype(1.0),
-                wptype(2.0),
-                wptype(3.0),
-                wptype(4.0),
-                wptype(5.0),
-                wptype(6.0),
-                wptype(7.0),
-                wptype(8.0),
-                wptype(9.0),
-                wptype(10.0),
-                wptype(11.0),
-                wptype(12.0),
-                wptype(13.0),
-                wptype(14.0),
-                wptype(15.0),
-                wptype(16.0),
-            ),
-        )
-
-        wp.expect_neq(
-            mat44(
-                wptype(1.0),
-                wptype(2.0),
-                wptype(3.0),
-                wptype(4.0),
-                wptype(5.0),
-                wptype(6.0),
-                wptype(7.0),
-                wptype(8.0),
-                wptype(9.0),
-                wptype(10.0),
-                wptype(11.0),
-                wptype(12.0),
-                wptype(13.0),
-                wptype(14.0),
-                wptype(15.0),
-                wptype(16.0),
-            ),
-            mat44(
-                -wptype(1.0),
-                wptype(2.0),
-                wptype(3.0),
-                wptype(4.0),
-                wptype(5.0),
-                wptype(6.0),
-                wptype(7.0),
-                wptype(8.0),
-                wptype(9.0),
-                wptype(10.0),
-                wptype(11.0),
-                wptype(12.0),
-                wptype(13.0),
-                wptype(14.0),
-                wptype(15.0),
-                wptype(16.0),
-            ),
-        )
-
-        wp.expect_eq(
-            mat55(
-                wptype(1.0),
-                wptype(2.0),
-                wptype(3.0),
-                wptype(4.0),
-                wptype(5.0),
-                wptype(6.0),
-                wptype(7.0),
-                wptype(8.0),
-                wptype(9.0),
-                wptype(10.0),
-                wptype(11.0),
-                wptype(12.0),
-                wptype(13.0),
-                wptype(14.0),
-                wptype(15.0),
-                wptype(16.0),
-                wptype(17.0),
-                wptype(18.0),
-                wptype(19.0),
-                wptype(20.0),
-                wptype(21.0),
-                wptype(22.0),
-                wptype(23.0),
-                wptype(24.0),
-                wptype(25.0),
-            ),
-            mat55(
-                wptype(1.0),
-                wptype(2.0),
-                wptype(3.0),
-                wptype(4.0),
-                wptype(5.0),
-                wptype(6.0),
-                wptype(7.0),
-                wptype(8.0),
-                wptype(9.0),
-                wptype(10.0),
-                wptype(11.0),
-                wptype(12.0),
-                wptype(13.0),
-                wptype(14.0),
-                wptype(15.0),
-                wptype(16.0),
-                wptype(17.0),
-                wptype(18.0),
-                wptype(19.0),
-                wptype(20.0),
-                wptype(21.0),
-                wptype(22.0),
-                wptype(23.0),
-                wptype(24.0),
-                wptype(25.0),
-            ),
-        )
-
-        wp.expect_neq(
-            mat55(
-                wptype(1.0),
-                wptype(2.0),
-                wptype(3.0),
-                wptype(4.0),
-                wptype(5.0),
-                wptype(6.0),
-                wptype(7.0),
-                wptype(8.0),
-                wptype(9.0),
-                wptype(10.0),
-                wptype(11.0),
-                wptype(12.0),
-                wptype(13.0),
-                wptype(14.0),
-                wptype(15.0),
-                wptype(16.0),
-                wptype(17.0),
-                wptype(18.0),
-                wptype(19.0),
-                wptype(20.0),
-                wptype(21.0),
-                wptype(22.0),
-                wptype(23.0),
-                wptype(24.0),
-                wptype(25.0),
-            ),
-            mat55(
-                wptype(1.0),
-                wptype(2.0),
-                wptype(3.0),
-                wptype(4.0),
-                wptype(5.0),
-                wptype(6.0),
-                wptype(7.0),
-                wptype(8.0),
-                wptype(9.0),
-                wptype(10.0),
-                wptype(11.0),
-                wptype(12.0),
-                wptype(13.0),
-                wptype(14.0),
-                wptype(15.0),
-                wptype(16.0),
-                -wptype(17.0),
-                wptype(18.0),
-                wptype(19.0),
-                wptype(20.0),
-                wptype(21.0),
-                wptype(22.0),
-                wptype(23.0),
-                wptype(24.0),
-                wptype(25.0),
-            ),
-        )
-
-    kernel = getkernel(check_mat_equality, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    wp.launch(kernel, dim=1, inputs=[], outputs=[], device=device)
-
-
-def test_negation(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 1.0e-2,
-        np.float32: 1.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-
-    output_select_kernel = get_select_kernel(wptype)
-
-    def check_mat_negation(
-        m2: wp.array(dtype=mat22),
-        m3: wp.array(dtype=mat33),
-        m4: wp.array(dtype=mat44),
-        m5: wp.array(dtype=mat55),
-        outcomponents: wp.array(dtype=wptype),
-    ):
-        mat2 = -m2[0]
-        mat3 = -m3[0]
-        mat4 = -m4[0]
-        mat5 = -m5[0]
-
-        # multiply outputs by 2 so we've got something to backpropagate:
-        idx = 0
-        for i in range(2):
-            for j in range(2):
-                outcomponents[idx] = wptype(2) * mat2[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(3):
-                outcomponents[idx] = wptype(2) * mat3[i, j]
-                idx = idx + 1
-
-        for i in range(4):
-            for j in range(4):
-                outcomponents[idx] = wptype(2) * mat4[i, j]
-                idx = idx + 1
-
-        for i in range(5):
-            for j in range(5):
-                outcomponents[idx] = wptype(2) * mat5[i, j]
-                idx = idx + 1
-
-    kernel = getkernel(check_mat_negation, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    m2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    m3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    m4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    m5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    outcomponents = wp.zeros(2 * 2 + 3 * 3 + 4 * 4 + 5 * 5, dtype=wptype, requires_grad=True, device=device)
-
-    wp.launch(kernel, dim=1, inputs=[m2, m3, m4, m5], outputs=[outcomponents], device=device)
-
-    assert_np_equal(outcomponents.numpy()[:4], -2 * m2.numpy().reshape(-1), tol=tol)
-    assert_np_equal(outcomponents.numpy()[4:13], -2 * m3.numpy().reshape(-1), tol=tol)
-    assert_np_equal(outcomponents.numpy()[13:29], -2 * m4.numpy().reshape(-1), tol=tol)
-    assert_np_equal(outcomponents.numpy()[29:54], -2 * m5.numpy().reshape(-1), tol=tol)
-
-    if dtype in np_float_types:
-        idx = 0
-        out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-        for dim, input in [(2, m2), (3, m3), (4, m4), (5, m5)]:
-            for i in range(dim):
-                for j in range(dim):
-                    tape = wp.Tape()
-                    with tape:
-                        wp.launch(kernel, dim=1, inputs=[m2, m3, m4, m5], outputs=[outcomponents], device=device)
-                        wp.launch(
-                            output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device
-                        )
-                    tape.backward(loss=out)
-                    expectedresult = np.zeros((dim, dim), dtype=dtype)
-                    expectedresult[i, j] = -2
-                    assert_np_equal(tape.gradients[input].numpy()[0], expectedresult)
-                    tape.zero()
-                    idx = idx + 1
-
-
-def test_transpose(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 1.0e-2,
-        np.float32: 1.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat32 = wp.types.matrix(shape=(3, 2), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-
-    output_select_kernel = get_select_kernel(wptype)
-
-    def check_mat_transpose(
-        m2: wp.array(dtype=mat22),
-        m3: wp.array(dtype=mat33),
-        m4: wp.array(dtype=mat44),
-        m5: wp.array(dtype=mat55),
-        m32: wp.array(dtype=mat32),
-        outcomponents: wp.array(dtype=wptype),
-    ):
-        # multiply outputs by 2 so we've got something to backpropagate:
-        mat2 = wptype(2) * wp.transpose(m2[0])
-        mat3 = wptype(2) * wp.transpose(m3[0])
-        mat4 = wptype(2) * wp.transpose(m4[0])
-        mat5 = wptype(2) * wp.transpose(m5[0])
-        mat32 = wptype(2) * wp.transpose(m32[0])
-
-        idx = 0
-        for i in range(2):
-            for j in range(2):
-                outcomponents[idx] = mat2[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(3):
-                outcomponents[idx] = mat3[i, j]
-                idx = idx + 1
-
-        for i in range(4):
-            for j in range(4):
-                outcomponents[idx] = mat4[i, j]
-                idx = idx + 1
-
-        for i in range(5):
-            for j in range(5):
-                outcomponents[idx] = mat5[i, j]
-                idx = idx + 1
-
-        for i in range(2):
-            for j in range(3):
-                outcomponents[idx] = mat32[i, j]
-                idx = idx + 1
-
-    kernel = getkernel(check_mat_transpose, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    m2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    m3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    m4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    m5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    m32 = wp.array(randvals([1, 3, 2], dtype), dtype=mat32, requires_grad=True, device=device)
-    outcomponents = wp.zeros(2 * 2 + 3 * 3 + 4 * 4 + 5 * 5 + 2 * 3, dtype=wptype, requires_grad=True, device=device)
-
-    wp.launch(kernel, dim=1, inputs=[m2, m3, m4, m5, m32], outputs=[outcomponents], device=device)
-
-    assert_np_equal(outcomponents.numpy()[:4], 2 * m2.numpy()[0].T.reshape(-1), tol=tol)
-    assert_np_equal(outcomponents.numpy()[4:13], 2 * m3.numpy()[0].T.reshape(-1), tol=tol)
-    assert_np_equal(outcomponents.numpy()[13:29], 2 * m4.numpy()[0].T.reshape(-1), tol=tol)
-    assert_np_equal(outcomponents.numpy()[29:54], 2 * m5.numpy()[0].T.reshape(-1), tol=tol)
-    assert_np_equal(outcomponents.numpy()[54:], 2 * m32.numpy()[0].T.reshape(-1), tol=tol)
-
-    if dtype in np_float_types:
-        idx = 0
-        out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-        for input in [m2, m3, m4, m5]:
-            for i in range(input.dtype._shape_[0]):
-                for j in range(input.dtype._shape_[1]):
-                    tape = wp.Tape()
-                    with tape:
-                        wp.launch(kernel, dim=1, inputs=[m2, m3, m4, m5, m32], outputs=[outcomponents], device=device)
-                        wp.launch(
-                            output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device
-                        )
-                    tape.backward(loss=out)
-                    expectedresult = np.zeros((input.dtype._shape_[1], input.dtype._shape_[0]), dtype=dtype)
-                    expectedresult[j, i] = 2
-                    assert_np_equal(tape.gradients[input].numpy()[0], expectedresult)
-                    tape.zero()
-                    idx = idx + 1
-
-
-def test_scalar_multiplication(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 1.0e-2,
-        np.float32: 1.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-
-    output_select_kernel = get_select_kernel(wptype)
-
-    def check_mat_scalar_mul(
-        s: wp.array(dtype=wptype),
-        m2: wp.array(dtype=mat22),
-        m3: wp.array(dtype=mat33),
-        m4: wp.array(dtype=mat44),
-        m5: wp.array(dtype=mat55),
-        outcomponents: wp.array(dtype=wptype),
-        outcomponents_rightmul: wp.array(dtype=wptype),
-    ):
-        m2result = s[0] * m2[0]
-        m3result = s[0] * m3[0]
-        m4result = s[0] * m4[0]
-        m5result = s[0] * m5[0]
-
-        m2resultright = m2[0] * s[0]
-        m3resultright = m3[0] * s[0]
-        m4resultright = m4[0] * s[0]
-        m5resultright = m5[0] * s[0]
-
-        m2result_2 = s[0] * m2[0]
-        m3result_2 = s[0] * m3[0]
-        m4result_2 = s[0] * m4[0]
-        m5result_2 = s[0] * m5[0]
-
-        m2resultright_2 = m2[0] * s[0]
-        m3resultright_2 = m3[0] * s[0]
-        m4resultright_2 = m4[0] * s[0]
-        m5resultright_2 = m5[0] * s[0]
-
-        # multiply outputs by 2 so we've got something to backpropagate:
-        idx = 0
-        for i in range(2):
-            for j in range(2):
-                outcomponents[idx] = wptype(2) * m2result[i, j]
-                outcomponents_rightmul[idx] = wptype(2) * m2resultright[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(3):
-                outcomponents[idx] = wptype(2) * m3result[i, j]
-                outcomponents_rightmul[idx] = wptype(2) * m3resultright[i, j]
-                idx = idx + 1
-
-        for i in range(4):
-            for j in range(4):
-                outcomponents[idx] = wptype(2) * m4result[i, j]
-                outcomponents_rightmul[idx] = wptype(2) * m4resultright[i, j]
-                idx = idx + 1
-
-        for i in range(5):
-            for j in range(5):
-                outcomponents[idx] = wptype(2) * m5result[i, j]
-                outcomponents_rightmul[idx] = wptype(2) * m5resultright[i, j]
-                idx = idx + 1
-
-        for i in range(2):
-            for j in range(2):
-                outcomponents[idx] = wptype(2) * m2result_2[i, j]
-                outcomponents_rightmul[idx] = wptype(2) * m2resultright_2[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(3):
-                outcomponents[idx] = wptype(2) * m3result_2[i, j]
-                outcomponents_rightmul[idx] = wptype(2) * m3resultright_2[i, j]
-                idx = idx + 1
-
-        for i in range(4):
-            for j in range(4):
-                outcomponents[idx] = wptype(2) * m4result_2[i, j]
-                outcomponents_rightmul[idx] = wptype(2) * m4resultright_2[i, j]
-                idx = idx + 1
-
-        for i in range(5):
-            for j in range(5):
-                outcomponents[idx] = wptype(2) * m5result_2[i, j]
-                outcomponents_rightmul[idx] = wptype(2) * m5resultright_2[i, j]
-                idx = idx + 1
-
-    kernel = getkernel(check_mat_scalar_mul, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    s = wp.array(randvals([1], dtype), requires_grad=True, device=device)
-    m2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    m3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    m4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    m5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    outcomponents = wp.zeros(2 * (2 * 2 + 3 * 3 + 4 * 4 + 5 * 5), dtype=wptype, requires_grad=True, device=device)
-    outcomponents_rightmul = wp.zeros(
-        2 * (2 * 2 + 3 * 3 + 4 * 4 + 5 * 5), dtype=wptype, requires_grad=True, device=device
-    )
-
-    wp.launch(kernel, dim=1, inputs=[s, m2, m3, m4, m5], outputs=[outcomponents, outcomponents_rightmul], device=device)
-
-    sval = s.numpy()[0]
-    assert_np_equal(outcomponents.numpy()[:4], 2 * sval * m2.numpy().reshape(-1), tol=tol)
-    assert_np_equal(outcomponents.numpy()[4:13], 2 * sval * m3.numpy().reshape(-1), tol=10 * tol)
-    assert_np_equal(outcomponents.numpy()[13:29], 2 * sval * m4.numpy().reshape(-1), tol=10 * tol)
-    assert_np_equal(outcomponents.numpy()[29:54], 2 * sval * m5.numpy().reshape(-1), tol=10 * tol)
-
-    assert_np_equal(outcomponents_rightmul.numpy()[:4], 2 * sval * m2.numpy().reshape(-1), tol=tol)
-    assert_np_equal(outcomponents_rightmul.numpy()[4:13], 2 * sval * m3.numpy().reshape(-1), tol=10 * tol)
-    assert_np_equal(outcomponents_rightmul.numpy()[13:29], 2 * sval * m4.numpy().reshape(-1), tol=10 * tol)
-    assert_np_equal(outcomponents_rightmul.numpy()[29:54], 2 * sval * m5.numpy().reshape(-1), tol=10 * tol)
-
-    assert_np_equal(outcomponents.numpy()[54:58], 2 * sval * m2.numpy().reshape(-1), tol=tol)
-    assert_np_equal(outcomponents.numpy()[58:67], 2 * sval * m3.numpy().reshape(-1), tol=10 * tol)
-    assert_np_equal(outcomponents.numpy()[67:83], 2 * sval * m4.numpy().reshape(-1), tol=10 * tol)
-    assert_np_equal(outcomponents.numpy()[83:108], 2 * sval * m5.numpy().reshape(-1), tol=10 * tol)
-
-    assert_np_equal(outcomponents_rightmul.numpy()[54:58], 2 * sval * m2.numpy().reshape(-1), tol=tol)
-    assert_np_equal(outcomponents_rightmul.numpy()[58:67], 2 * sval * m3.numpy().reshape(-1), tol=10 * tol)
-    assert_np_equal(outcomponents_rightmul.numpy()[67:83], 2 * sval * m4.numpy().reshape(-1), tol=10 * tol)
-    assert_np_equal(outcomponents_rightmul.numpy()[83:108], 2 * sval * m5.numpy().reshape(-1), tol=10 * tol)
-
-    if dtype in np_float_types:
-        idx = 0
-        out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-        for dim, input in [(2, m2), (3, m3), (4, m4), (5, m5)]:
-            for i in range(dim):
-                for j in range(dim):
-                    # test left mul gradient:
-                    tape = wp.Tape()
-                    with tape:
-                        wp.launch(
-                            kernel,
-                            dim=1,
-                            inputs=[s, m2, m3, m4, m5],
-                            outputs=[outcomponents, outcomponents_rightmul],
-                            device=device,
-                        )
-                        wp.launch(
-                            output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device
-                        )
-                    tape.backward(loss=out)
-                    expectedresult = np.zeros((dim, dim), dtype=dtype)
-                    expectedresult[i, j] = 2 * sval
-                    assert_np_equal(tape.gradients[input].numpy()[0], expectedresult, tol=10 * tol)
-                    assert_np_equal(tape.gradients[s].numpy()[0], 2 * input.numpy()[0, i, j], tol=10 * tol)
-                    tape.zero()
-
-                    # test right mul gradient:
-                    tape = wp.Tape()
-                    with tape:
-                        wp.launch(
-                            kernel,
-                            dim=1,
-                            inputs=[s, m2, m3, m4, m5],
-                            outputs=[outcomponents, outcomponents_rightmul],
-                            device=device,
-                        )
-                        wp.launch(
-                            output_select_kernel,
-                            dim=1,
-                            inputs=[outcomponents_rightmul, idx],
-                            outputs=[out],
-                            device=device,
-                        )
-                    tape.backward(loss=out)
-                    expectedresult = np.zeros((dim, dim), dtype=dtype)
-                    expectedresult[i, j] = 2 * sval
-                    assert_np_equal(tape.gradients[input].numpy()[0], expectedresult, tol=10 * tol)
-                    assert_np_equal(tape.gradients[s].numpy()[0], 2 * input.numpy()[0, i, j], tol=10 * tol)
-                    tape.zero()
-
-                    idx = idx + 1
-
-
-def test_matvec_multiplication(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 2.0e-2,
-        np.float32: 5.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat32 = wp.types.matrix(shape=(3, 2), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-
-    vec2 = wp.types.vector(length=2, dtype=wptype)
-    vec3 = wp.types.vector(length=3, dtype=wptype)
-    vec4 = wp.types.vector(length=4, dtype=wptype)
-    vec5 = wp.types.vector(length=5, dtype=wptype)
-
-    output_select_kernel = get_select_kernel(wptype)
-
-    def check_mat_vec_mul(
-        v2: wp.array(dtype=vec2),
-        v3: wp.array(dtype=vec3),
-        v4: wp.array(dtype=vec4),
-        v5: wp.array(dtype=vec5),
-        v32: wp.array(dtype=vec2),
-        m2: wp.array(dtype=mat22),
-        m3: wp.array(dtype=mat33),
-        m4: wp.array(dtype=mat44),
-        m5: wp.array(dtype=mat55),
-        m32: wp.array(dtype=mat32),
-        outcomponents: wp.array(dtype=wptype),
-    ):
-        v2result = m2[0] * v2[0]
-        v3result = m3[0] * v3[0]
-        v4result = m4[0] * v4[0]
-        v5result = m5[0] * v5[0]
-        v32result = m32[0] * v32[0]
-        v2result_2 = m2[0] @ v2[0]
-        v3result_2 = m3[0] @ v3[0]
-        v4result_2 = m4[0] @ v4[0]
-        v5result_2 = m5[0] @ v5[0]
-        v32result_2 = m32[0] @ v32[0]
-
-        idx = 0
-
-        # multiply outputs by 2 so we've got something to backpropagate:
-        for i in range(2):
-            outcomponents[idx] = wptype(2) * v2result[i]
-            idx = idx + 1
-
-        for i in range(3):
-            outcomponents[idx] = wptype(2) * v3result[i]
-            idx = idx + 1
-
-        for i in range(4):
-            outcomponents[idx] = wptype(2) * v4result[i]
-            idx = idx + 1
-
-        for i in range(5):
-            outcomponents[idx] = wptype(2) * v5result[i]
-            idx = idx + 1
-
-        for i in range(3):
-            outcomponents[idx] = wptype(2) * v32result[i]
-            idx = idx + 1
-
-        for i in range(2):
-            outcomponents[idx] = wptype(2) * v2result_2[i]
-            idx = idx + 1
-
-        for i in range(3):
-            outcomponents[idx] = wptype(2) * v3result_2[i]
-            idx = idx + 1
-
-        for i in range(4):
-            outcomponents[idx] = wptype(2) * v4result_2[i]
-            idx = idx + 1
-
-        for i in range(5):
-            outcomponents[idx] = wptype(2) * v5result_2[i]
-            idx = idx + 1
-
-        for i in range(3):
-            outcomponents[idx] = wptype(2) * v32result_2[i]
-            idx = idx + 1
-
-    kernel = getkernel(check_mat_vec_mul, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    v2 = wp.array(randvals([1, 2], dtype), dtype=vec2, requires_grad=True, device=device)
-    v3 = wp.array(randvals([1, 3], dtype), dtype=vec3, requires_grad=True, device=device)
-    v4 = wp.array(randvals([1, 4], dtype), dtype=vec4, requires_grad=True, device=device)
-    v5 = wp.array(randvals([1, 5], dtype), dtype=vec5, requires_grad=True, device=device)
-    v32 = wp.array(randvals([1, 2], dtype), dtype=vec2, requires_grad=True, device=device)
-    m2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    m3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    m4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    m5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    m32 = wp.array(randvals([1, 3, 2], dtype), dtype=mat32, requires_grad=True, device=device)
-    outcomponents = wp.zeros(2 * (2 + 3 + 4 + 5 + 3), dtype=wptype, requires_grad=True, device=device)
-
-    wp.launch(kernel, dim=1, inputs=[v2, v3, v4, v5, v32, m2, m3, m4, m5, m32], outputs=[outcomponents], device=device)
-
-    assert_np_equal(outcomponents.numpy()[:2], 2 * np.matmul(m2.numpy()[0], v2.numpy()[0]), tol=tol)
-    assert_np_equal(outcomponents.numpy()[2:5], 2 * np.matmul(m3.numpy()[0], v3.numpy()[0]), tol=tol)
-    assert_np_equal(outcomponents.numpy()[5:9], 2 * np.matmul(m4.numpy()[0], v4.numpy()[0]), tol=5 * tol)
-    assert_np_equal(outcomponents.numpy()[9:14], 2 * np.matmul(m5.numpy()[0], v5.numpy()[0]), tol=5 * tol)
-    assert_np_equal(outcomponents.numpy()[14:17], 2 * np.matmul(m32.numpy()[0], v32.numpy()[0]), tol=5 * tol)
-    assert_np_equal(outcomponents.numpy()[17:19], 2 * np.matmul(m2.numpy()[0], v2.numpy()[0]), tol=tol)
-    assert_np_equal(outcomponents.numpy()[19:22], 2 * np.matmul(m3.numpy()[0], v3.numpy()[0]), tol=tol)
-    assert_np_equal(outcomponents.numpy()[22:26], 2 * np.matmul(m4.numpy()[0], v4.numpy()[0]), tol=5 * tol)
-    assert_np_equal(outcomponents.numpy()[26:31], 2 * np.matmul(m5.numpy()[0], v5.numpy()[0]), tol=5 * tol)
-    assert_np_equal(outcomponents.numpy()[31:34], 2 * np.matmul(m32.numpy()[0], v32.numpy()[0]), tol=5 * tol)
-
-    if dtype in np_float_types:
-        idx = 0
-        out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-        for dim, invec, inmat in [(2, v2, m2), (3, v3, m3), (4, v4, m4), (5, v5, m5), (3, v32, m32)]:
-            for i in range(dim):
-                tape = wp.Tape()
-                with tape:
-                    wp.launch(
-                        kernel,
-                        dim=1,
-                        inputs=[v2, v3, v4, v5, v32, m2, m3, m4, m5, m32],
-                        outputs=[outcomponents],
-                        device=device,
-                    )
-                    wp.launch(output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device)
-                tape.backward(loss=out)
-
-                assert_np_equal(tape.gradients[invec].numpy()[0], 2 * inmat.numpy()[0, i, :], tol=2 * tol)
-                expectedresult = np.zeros(inmat.dtype._shape_, dtype=dtype)
-                expectedresult[i, :] = 2 * invec.numpy()[0]
-                assert_np_equal(tape.gradients[inmat].numpy()[0], expectedresult, tol=2 * tol)
-
-                tape.zero()
-
-                idx = idx + 1
-
-
-def test_matmat_multiplication(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 2.0e-2,
-        np.float32: 5.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat32 = wp.types.matrix(shape=(3, 2), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-
-    output_select_kernel = get_select_kernel(wptype)
-
-    def check_mat_mat_mul(
-        a2: wp.array(dtype=mat22),
-        a3: wp.array(dtype=mat33),
-        a4: wp.array(dtype=mat44),
-        a5: wp.array(dtype=mat55),
-        a32: wp.array(dtype=mat32),
-        b2: wp.array(dtype=mat22),
-        b3: wp.array(dtype=mat33),
-        b4: wp.array(dtype=mat44),
-        b5: wp.array(dtype=mat55),
-        b32: wp.array(dtype=mat32),
-        outcomponents: wp.array(dtype=wptype),
-    ):
-        c2result = b2[0] * a2[0]
-        c3result = b3[0] * a3[0]
-        c4result = b4[0] * a4[0]
-        c5result = b5[0] * a5[0]
-        c32result = b32[0] * a2[0]
-        c32result2 = b3[0] * a32[0]
-        c2result_2 = b2[0] @ a2[0]
-        c3result_2 = b3[0] @ a3[0]
-        c4result_2 = b4[0] @ a4[0]
-        c5result_2 = b5[0] @ a5[0]
-        c32result_2 = b32[0] @ a2[0]
-        c32result2_2 = b3[0] @ a32[0]
-
-        # multiply outputs by 2 so we've got something to backpropagate:
-        idx = 0
-        for i in range(2):
-            for j in range(2):
-                outcomponents[idx] = wptype(2) * c2result[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(3):
-                outcomponents[idx] = wptype(2) * c3result[i, j]
-                idx = idx + 1
-
-        for i in range(4):
-            for j in range(4):
-                outcomponents[idx] = wptype(2) * c4result[i, j]
-                idx = idx + 1
-
-        for i in range(5):
-            for j in range(5):
-                outcomponents[idx] = wptype(2) * c5result[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(2):
-                outcomponents[idx] = wptype(2) * c32result[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(2):
-                outcomponents[idx] = wptype(2) * c32result2[i, j]
-                idx = idx + 1
-
-        for i in range(2):
-            for j in range(2):
-                outcomponents[idx] = wptype(2) * c2result_2[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(3):
-                outcomponents[idx] = wptype(2) * c3result_2[i, j]
-                idx = idx + 1
-
-        for i in range(4):
-            for j in range(4):
-                outcomponents[idx] = wptype(2) * c4result_2[i, j]
-                idx = idx + 1
-
-        for i in range(5):
-            for j in range(5):
-                outcomponents[idx] = wptype(2) * c5result_2[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(2):
-                outcomponents[idx] = wptype(2) * c32result_2[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(2):
-                outcomponents[idx] = wptype(2) * c32result2_2[i, j]
-                idx = idx + 1
-
-    kernel = getkernel(check_mat_mat_mul, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    v2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    v3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    v4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    v5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    v32 = wp.array(randvals([1, 3, 2], dtype), dtype=mat32, requires_grad=True, device=device)
-    m2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    m3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    m4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    m5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    m32 = wp.array(randvals([1, 3, 2], dtype), dtype=mat32, requires_grad=True, device=device)
-    outcomponents = wp.zeros(
-        2 * (2 * 2 + 3 * 3 + 4 * 4 + 5 * 5 + 3 * 2 + 3 * 2), dtype=wptype, requires_grad=True, device=device
-    )
-
-    wp.launch(kernel, dim=1, inputs=[v2, v3, v4, v5, v32, m2, m3, m4, m5, m32], outputs=[outcomponents], device=device)
-
-    assert_np_equal(outcomponents.numpy()[:4], 2 * np.matmul(m2.numpy()[0], v2.numpy()[0]), tol=tol)
-    assert_np_equal(outcomponents.numpy()[4:13], 2 * np.matmul(m3.numpy()[0], v3.numpy()[0]), tol=tol)
-    assert_np_equal(outcomponents.numpy()[13:29], 2 * np.matmul(m4.numpy()[0], v4.numpy()[0]), tol=2 * tol)
-    assert_np_equal(outcomponents.numpy()[29:54], 2 * np.matmul(m5.numpy()[0], v5.numpy()[0]), tol=10 * tol)
-    assert_np_equal(outcomponents.numpy()[54:60], 2 * np.matmul(m32.numpy()[0], v2.numpy()[0]), tol=5 * tol)
-    assert_np_equal(outcomponents.numpy()[60:66], 2 * np.matmul(m3.numpy()[0], v32.numpy()[0]), tol=5 * tol)
-    assert_np_equal(outcomponents.numpy()[66:70], 2 * np.matmul(m2.numpy()[0], v2.numpy()[0]), tol=tol)
-    assert_np_equal(outcomponents.numpy()[70:79], 2 * np.matmul(m3.numpy()[0], v3.numpy()[0]), tol=tol)
-    assert_np_equal(outcomponents.numpy()[79:95], 2 * np.matmul(m4.numpy()[0], v4.numpy()[0]), tol=2 * tol)
-    assert_np_equal(outcomponents.numpy()[95:120], 2 * np.matmul(m5.numpy()[0], v5.numpy()[0]), tol=10 * tol)
-    assert_np_equal(outcomponents.numpy()[120:126], 2 * np.matmul(m32.numpy()[0], v2.numpy()[0]), tol=5 * tol)
-    assert_np_equal(outcomponents.numpy()[126:132], 2 * np.matmul(m3.numpy()[0], v32.numpy()[0]), tol=5 * tol)
-
-    if dtype in np_float_types:
-        idx = 0
-        out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-        for v, m in [(v2, m2), (v3, m3), (v4, m4), (v5, m5), (v2, m32), (v32, m3)]:
-            rows, cols = m.dtype._shape_[0], v.dtype._shape_[1]
-            for i in range(rows):
-                for j in range(cols):
-                    tape = wp.Tape()
-                    with tape:
-                        wp.launch(
-                            kernel,
-                            dim=1,
-                            inputs=[v2, v3, v4, v5, v32, m2, m3, m4, m5, m32],
-                            outputs=[outcomponents],
-                            device=device,
-                        )
-                        wp.launch(
-                            output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device
-                        )
-                    tape.backward(loss=out)
-
-                    expected = np.zeros(v.dtype._shape_, dtype=dtype)
-                    expected[:, j] = 2 * m.numpy()[0, i, :]
-                    assert_np_equal(tape.gradients[v].numpy()[0], expected, tol=10 * tol)
-
-                    expected = np.zeros(m.dtype._shape_, dtype=dtype)
-                    expected[i, :] = 2 * v.numpy()[0, :, j]
-                    assert_np_equal(tape.gradients[m].numpy()[0], expected, tol=10 * tol)
-
-                    tape.zero()
-                    idx = idx + 1
-
-
-def test_cw_multiplication(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 5.0e-2,
-        np.float32: 1.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-
-    output_select_kernel = get_select_kernel(wptype)
-
-    def check_mat_cw_mul(
-        s2: wp.array(dtype=mat22),
-        s3: wp.array(dtype=mat33),
-        s4: wp.array(dtype=mat44),
-        s5: wp.array(dtype=mat55),
-        v2: wp.array(dtype=mat22),
-        v3: wp.array(dtype=mat33),
-        v4: wp.array(dtype=mat44),
-        v5: wp.array(dtype=mat55),
-        outcomponents: wp.array(dtype=wptype),
-    ):
-        v2result = wptype(2) * wp.cw_mul(v2[0], s2[0])
-        v3result = wptype(2) * wp.cw_mul(v3[0], s3[0])
-        v4result = wptype(2) * wp.cw_mul(v4[0], s4[0])
-        v5result = wptype(2) * wp.cw_mul(v5[0], s5[0])
-
-        # multiply outputs by 2 so we've got something to backpropagate:
-        idx = 0
-        for i in range(2):
-            for j in range(2):
-                outcomponents[idx] = v2result[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(3):
-                outcomponents[idx] = v3result[i, j]
-                idx = idx + 1
-
-        for i in range(4):
-            for j in range(4):
-                outcomponents[idx] = v4result[i, j]
-                idx = idx + 1
-
-        for i in range(5):
-            for j in range(5):
-                outcomponents[idx] = v5result[i, j]
-                idx = idx + 1
-
-    kernel = getkernel(check_mat_cw_mul, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    s2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    s3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    s4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    s5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    v2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    v3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    v4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    v5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    outcomponents = wp.zeros(2 * 2 + 3 * 3 + 4 * 4 + 5 * 5, dtype=wptype, requires_grad=True, device=device)
-
-    wp.launch(
-        kernel,
-        dim=1,
-        inputs=[
-            s2,
-            s3,
-            s4,
-            s5,
-            v2,
-            v3,
-            v4,
-            v5,
-        ],
-        outputs=[outcomponents],
-        device=device,
-    )
-
-    assert_np_equal(outcomponents.numpy()[:4], 2 * (v2.numpy() * s2.numpy()).reshape(-1), tol=50 * tol)
-    assert_np_equal(outcomponents.numpy()[4:13], 2 * (v3.numpy() * s3.numpy()).reshape(-1), tol=50 * tol)
-    assert_np_equal(outcomponents.numpy()[13:29], 2 * (v4.numpy() * s4.numpy()).reshape(-1), tol=50 * tol)
-    assert_np_equal(outcomponents.numpy()[29:54], 2 * (v5.numpy() * s5.numpy()).reshape(-1), tol=50 * tol)
-
-    if dtype in np_float_types:
-        idx = 0
-        out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-        for dim, in1, in2 in [(2, s2, v2), (3, s3, v3), (4, s4, v4), (5, s5, v5)]:
-            for i in range(dim):
-                for j in range(dim):
-                    tape = wp.Tape()
-                    with tape:
-                        wp.launch(
-                            kernel,
-                            dim=1,
-                            inputs=[
-                                s2,
-                                s3,
-                                s4,
-                                s5,
-                                v2,
-                                v3,
-                                v4,
-                                v5,
-                            ],
-                            outputs=[outcomponents],
-                            device=device,
-                        )
-                        wp.launch(
-                            output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device
-                        )
-                    tape.backward(loss=out)
-                    expectedresult = np.zeros((dim, dim), dtype=dtype)
-                    expectedresult[i, j] = 2 * in1.numpy()[0][i, j]
-                    assert_np_equal(tape.gradients[in2].numpy()[0], expectedresult, tol=5 * tol)
-                    expectedresult[i, j] = 2 * in2.numpy()[0][i, j]
-                    assert_np_equal(tape.gradients[in1].numpy()[0], expectedresult, tol=5 * tol)
-                    tape.zero()
-
-                    idx = idx + 1
-
-
-def test_cw_division(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 1.0e-2,
-        np.float32: 1.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-
-    output_select_kernel = get_select_kernel(wptype)
-
-    def check_mat_cw_div(
-        s2: wp.array(dtype=mat22),
-        s3: wp.array(dtype=mat33),
-        s4: wp.array(dtype=mat44),
-        s5: wp.array(dtype=mat55),
-        v2: wp.array(dtype=mat22),
-        v3: wp.array(dtype=mat33),
-        v4: wp.array(dtype=mat44),
-        v5: wp.array(dtype=mat55),
-        outcomponents: wp.array(dtype=wptype),
-    ):
-        v2result = wptype(2) * wp.cw_div(v2[0], s2[0])
-        v3result = wptype(2) * wp.cw_div(v3[0], s3[0])
-        v4result = wptype(2) * wp.cw_div(v4[0], s4[0])
-        v5result = wptype(2) * wp.cw_div(v5[0], s5[0])
-
-        # multiply outputs by 2 so we've got something to backpropagate:
-        idx = 0
-        for i in range(2):
-            for j in range(2):
-                outcomponents[idx] = v2result[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(3):
-                outcomponents[idx] = v3result[i, j]
-                idx = idx + 1
-
-        for i in range(4):
-            for j in range(4):
-                outcomponents[idx] = v4result[i, j]
-                idx = idx + 1
-
-        for i in range(5):
-            for j in range(5):
-                outcomponents[idx] = v5result[i, j]
-                idx = idx + 1
-
-    kernel = getkernel(check_mat_cw_div, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    s2 = randvals([1, 2, 2], dtype)
-    s3 = randvals([1, 3, 3], dtype)
-    s4 = randvals([1, 4, 4], dtype)
-    s5 = randvals([1, 5, 5], dtype)
-
-    # set denominators to 1 if their magnitudes are small
-    # to prevent divide by zero, or overflows if we're testing
-    # float16:
-    s2[np.abs(s2) < 1.0e-2] = 1
-    s3[np.abs(s3) < 1.0e-2] = 1
-    s4[np.abs(s4) < 1.0e-2] = 1
-    s5[np.abs(s5) < 1.0e-2] = 1
-
-    s2 = wp.array(s2, dtype=mat22, requires_grad=True, device=device)
-    s3 = wp.array(s3, dtype=mat33, requires_grad=True, device=device)
-    s4 = wp.array(s4, dtype=mat44, requires_grad=True, device=device)
-    s5 = wp.array(s5, dtype=mat55, requires_grad=True, device=device)
-
-    v2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    v3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    v4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    v5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    outcomponents = wp.zeros(2 * 2 + 3 * 3 + 4 * 4 + 5 * 5, dtype=wptype, requires_grad=True, device=device)
-
-    wp.launch(
-        kernel,
-        dim=1,
-        inputs=[
-            s2,
-            s3,
-            s4,
-            s5,
-            v2,
-            v3,
-            v4,
-            v5,
-        ],
-        outputs=[outcomponents],
-        device=device,
-    )
-
-    if dtype in np_float_types:
-        assert_np_equal(outcomponents.numpy()[:4], 2 * (v2.numpy() / s2.numpy()).reshape(-1), tol=50 * tol)
-        assert_np_equal(outcomponents.numpy()[4:13], 2 * (v3.numpy() / s3.numpy()).reshape(-1), tol=50 * tol)
-        assert_np_equal(outcomponents.numpy()[13:29], 2 * (v4.numpy() / s4.numpy()).reshape(-1), tol=50 * tol)
-        assert_np_equal(outcomponents.numpy()[29:54], 2 * (v5.numpy() / s5.numpy()).reshape(-1), tol=50 * tol)
-    else:
-        assert_np_equal(outcomponents.numpy()[:4], 2 * (v2.numpy() // s2.numpy()).reshape(-1), tol=50 * tol)
-        assert_np_equal(outcomponents.numpy()[4:13], 2 * (v3.numpy() // s3.numpy()).reshape(-1), tol=50 * tol)
-        assert_np_equal(outcomponents.numpy()[13:29], 2 * (v4.numpy() // s4.numpy()).reshape(-1), tol=50 * tol)
-        assert_np_equal(outcomponents.numpy()[29:54], 2 * (v5.numpy() // s5.numpy()).reshape(-1), tol=50 * tol)
-
-    if dtype in np_float_types:
-        idx = 0
-        out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-        for dim, s, v in [(2, s2, v2), (3, s3, v3), (4, s4, v4), (5, s5, v5)]:
-            for i in range(dim):
-                for j in range(dim):
-                    tape = wp.Tape()
-                    with tape:
-                        wp.launch(
-                            kernel,
-                            dim=1,
-                            inputs=[
-                                s2,
-                                s3,
-                                s4,
-                                s5,
-                                v2,
-                                v3,
-                                v4,
-                                v5,
-                            ],
-                            outputs=[outcomponents],
-                            device=device,
-                        )
-                        wp.launch(
-                            output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device
-                        )
-                    tape.backward(loss=out)
-
-                    # y = v/s
-                    # dy/dv = 1.0/s
-                    # dy/ds = -v/s^2
-
-                    expectedresult = np.zeros((dim, dim), dtype=dtype)
-                    expectedresult[i, j] = 2.0 / (s.numpy()[0, i, j])
-                    assert_np_equal(tape.gradients[v].numpy()[0], expectedresult, tol=50 * tol)
-                    expectedresult[i, j] = -2.0 * v.numpy()[0, i, j] / (s.numpy()[0, i, j] ** 2)
-                    assert_np_equal(
-                        tape.gradients[s].numpy()[0], expectedresult, tol=abs(outcomponents.numpy()[idx]) * 50 * tol
-                    )
-                    tape.zero()
-
-                    idx = idx + 1
-
-
-def test_outer_product(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 5.0e-3,
-        np.float32: 1.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    vec2 = wp.types.vector(length=2, dtype=wptype)
-    vec3 = wp.types.vector(length=3, dtype=wptype)
-    vec4 = wp.types.vector(length=4, dtype=wptype)
-    vec5 = wp.types.vector(length=5, dtype=wptype)
-
-    output_select_kernel = get_select_kernel(wptype)
-
-    def check_mat_outer_product(
-        s2: wp.array(dtype=vec2),
-        s3: wp.array(dtype=vec3),
-        s4: wp.array(dtype=vec4),
-        s5: wp.array(dtype=vec5),
-        v2: wp.array(dtype=vec2),
-        v3: wp.array(dtype=vec3),
-        v4: wp.array(dtype=vec4),
-        v5: wp.array(dtype=vec5),
-        outcomponents: wp.array(dtype=wptype),
-    ):
-        m22result = wptype(2) * wp.outer(s2[0], v2[0])
-        m33result = wptype(2) * wp.outer(s3[0], v3[0])
-        m44result = wptype(2) * wp.outer(s4[0], v4[0])
-        m55result = wptype(2) * wp.outer(s5[0], v5[0])
-        m25result = wptype(2) * wp.outer(s2[0], v5[0])
-
-        # multiply outputs by 2 so we've got something to backpropagate:
-        idx = 0
-        for i in range(2):
-            for j in range(2):
-                outcomponents[idx] = m22result[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(3):
-                outcomponents[idx] = m33result[i, j]
-                idx = idx + 1
-
-        for i in range(4):
-            for j in range(4):
-                outcomponents[idx] = m44result[i, j]
-                idx = idx + 1
-
-        for i in range(5):
-            for j in range(5):
-                outcomponents[idx] = m55result[i, j]
-                idx = idx + 1
-
-        for i in range(2):
-            for j in range(5):
-                outcomponents[idx] = m25result[i, j]
-                idx = idx + 1
-
-    kernel = getkernel(check_mat_outer_product, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    s2 = wp.array(randvals([1, 2], dtype), dtype=vec2, requires_grad=True, device=device)
-    s3 = wp.array(randvals([1, 3], dtype), dtype=vec3, requires_grad=True, device=device)
-    s4 = wp.array(randvals([1, 4], dtype), dtype=vec4, requires_grad=True, device=device)
-    s5 = wp.array(randvals([1, 5], dtype), dtype=vec5, requires_grad=True, device=device)
-    v2 = wp.array(randvals([1, 2], dtype), dtype=vec2, requires_grad=True, device=device)
-    v3 = wp.array(randvals([1, 3], dtype), dtype=vec3, requires_grad=True, device=device)
-    v4 = wp.array(randvals([1, 4], dtype), dtype=vec4, requires_grad=True, device=device)
-    v5 = wp.array(randvals([1, 5], dtype), dtype=vec5, requires_grad=True, device=device)
-    outcomponents = wp.zeros(2 * 2 + 3 * 3 + 4 * 4 + 5 * 5 + 2 * 5, dtype=wptype, requires_grad=True, device=device)
-
-    wp.launch(kernel, dim=1, inputs=[s2, s3, s4, s5, v2, v3, v4, v5], outputs=[outcomponents], device=device)
-
-    assert_np_equal(outcomponents.numpy()[:4], 2 * s2.numpy()[0, :, None] * v2.numpy()[0, None, :], tol=tol)
-    assert_np_equal(outcomponents.numpy()[4:13], 2 * s3.numpy()[0, :, None] * v3.numpy()[0, None, :], tol=10 * tol)
-    assert_np_equal(outcomponents.numpy()[13:29], 2 * s4.numpy()[0, :, None] * v4.numpy()[0, None, :], tol=10 * tol)
-    assert_np_equal(outcomponents.numpy()[29:54], 2 * s5.numpy()[0, :, None] * v5.numpy()[0, None, :], tol=10 * tol)
-    assert_np_equal(outcomponents.numpy()[54:], 2 * s2.numpy()[0, :, None] * v5.numpy()[0, None, :], tol=10 * tol)
-
-    if dtype in np_float_types:
-        idx = 0
-        out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-        for s, v in [(s2, v2), (s3, v3), (s4, v4), (s5, v5), (s2, v5)]:
-            rows = s.dtype._length_
-            cols = v.dtype._length_
-            for i in range(rows):
-                for j in range(cols):
-                    tape = wp.Tape()
-                    with tape:
-                        wp.launch(
-                            kernel,
-                            dim=1,
-                            inputs=[
-                                s2,
-                                s3,
-                                s4,
-                                s5,
-                                v2,
-                                v3,
-                                v4,
-                                v5,
-                            ],
-                            outputs=[outcomponents],
-                            device=device,
-                        )
-                        wp.launch(
-                            output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device
-                        )
-                    tape.backward(loss=out)
-
-                    # this component's gonna be s_i * v_j, so its s gradient is gonna be nozero
-                    # at the ith component and its v gradient will be nonzero at the jth component:
-
-                    expectedresult = np.zeros((rows), dtype=dtype)
-                    expectedresult[i] = 2 * v.numpy()[0, j]
-                    assert_np_equal(tape.gradients[s].numpy()[0], expectedresult, tol=10 * tol)
-
-                    expectedresult = np.zeros((cols), dtype=dtype)
-                    expectedresult[j] = 2 * s.numpy()[0, i]
-                    assert_np_equal(tape.gradients[v].numpy()[0], expectedresult, tol=10 * tol)
-                    tape.zero()
-
-                    idx = idx + 1
-
-
-def test_scalar_division(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 1.0e-2,
-        np.float32: 1.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-
-    output_select_kernel = get_select_kernel(wptype)
-
-    def check_mat_scalar_div(
-        s: wp.array(dtype=wptype),
-        m2: wp.array(dtype=mat22),
-        m3: wp.array(dtype=mat33),
-        m4: wp.array(dtype=mat44),
-        m5: wp.array(dtype=mat55),
-        outcomponents: wp.array(dtype=wptype),
-    ):
-        m2result = m2[0] / s[0]
-        m3result = m3[0] / s[0]
-        m4result = m4[0] / s[0]
-        m5result = m5[0] / s[0]
-
-        # multiply outputs by 2 so we've got something to backpropagate:
-        idx = 0
-        for i in range(2):
-            for j in range(2):
-                outcomponents[idx] = wptype(2) * m2result[i, j]
-                idx = idx + 1
-
-        for i in range(3):
-            for j in range(3):
-                outcomponents[idx] = wptype(2) * m3result[i, j]
-                idx = idx + 1
-
-        for i in range(4):
-            for j in range(4):
-                outcomponents[idx] = wptype(2) * m4result[i, j]
-                idx = idx + 1
-
-        for i in range(5):
-            for j in range(5):
-                outcomponents[idx] = wptype(2) * m5result[i, j]
-                idx = idx + 1
-
-    kernel = getkernel(check_mat_scalar_div, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    s = wp.array(randvals([1], dtype), requires_grad=True, device=device)
-    m2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    m3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    m4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    m5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    outcomponents = wp.zeros(2 * 2 + 3 * 3 + 4 * 4 + 5 * 5, dtype=wptype, requires_grad=True, device=device)
-
-    wp.launch(kernel, dim=1, inputs=[s, m2, m3, m4, m5], outputs=[outcomponents], device=device)
-
-    sval = s.numpy()[0]
-    if dtype in np_float_types:
-        assert_np_equal(outcomponents.numpy()[:4], 2 * m2.numpy().reshape(-1) / sval, tol=tol)
-        assert_np_equal(outcomponents.numpy()[4:13], 2 * m3.numpy().reshape(-1) / sval, tol=10 * tol)
-        assert_np_equal(outcomponents.numpy()[13:29], 2 * m4.numpy().reshape(-1) / sval, tol=10 * tol)
-        assert_np_equal(outcomponents.numpy()[29:54], 2 * m5.numpy().reshape(-1) / sval, tol=10 * tol)
-    else:
-        assert_np_equal(outcomponents.numpy()[:4], 2 * (m2.numpy().reshape(-1) // sval), tol=tol)
-        assert_np_equal(outcomponents.numpy()[4:13], 2 * (m3.numpy().reshape(-1) // sval), tol=10 * tol)
-        assert_np_equal(outcomponents.numpy()[13:29], 2 * (m4.numpy().reshape(-1) // sval), tol=10 * tol)
-        assert_np_equal(outcomponents.numpy()[29:54], 2 * (m5.numpy().reshape(-1) // sval), tol=10 * tol)
-
-    if dtype in np_float_types:
-        idx = 0
-        out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-        for dim, input in [(2, m2), (3, m3), (4, m4), (5, m5)]:
-            for i in range(dim):
-                for j in range(dim):
-                    tape = wp.Tape()
-                    with tape:
-                        wp.launch(kernel, dim=1, inputs=[s, m2, m3, m4, m5], outputs=[outcomponents], device=device)
-                        wp.launch(
-                            output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device
-                        )
-                    tape.backward(loss=out)
-                    expectedresult = np.zeros((dim, dim), dtype=dtype)
-                    expectedresult[i, j] = 2.0 / sval
-                    assert_np_equal(tape.gradients[input].numpy()[0], expectedresult, tol=10 * tol)
-                    assert_np_equal(
-                        tape.gradients[s].numpy()[0], -2 * input.numpy()[0, i, j] / (sval * sval), tol=10 * tol
-                    )
-                    tape.zero()
-
-                    idx = idx + 1
-
-
-def test_addition(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 2.0e-2,
-        np.float32: 5.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-
-    output_select_kernel = get_select_kernel(wptype)
-
-    def check_mat_add(
-        s2: wp.array(dtype=mat22),
-        s3: wp.array(dtype=mat33),
-        s4: wp.array(dtype=mat44),
-        s5: wp.array(dtype=mat55),
-        v2: wp.array(dtype=mat22),
-        v3: wp.array(dtype=mat33),
-        v4: wp.array(dtype=mat44),
-        v5: wp.array(dtype=mat55),
+    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
+    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
+
+    output_select_kernel = get_select_kernel(wptype)
+
+    def check_mat_negation(
+        m2: wp.array(dtype=mat22),
+        m3: wp.array(dtype=mat33),
+        m4: wp.array(dtype=mat44),
+        m5: wp.array(dtype=mat55),
         outcomponents: wp.array(dtype=wptype),
     ):
-        v2result = v2[0] + s2[0]
-        v3result = v3[0] + s3[0]
-        v4result = v4[0] + s4[0]
-        v5result = v5[0] + s5[0]
+        mat2 = -m2[0]
+        mat3 = -m3[0]
+        mat4 = -m4[0]
+        mat5 = -m5[0]
 
         # multiply outputs by 2 so we've got something to backpropagate:
         idx = 0
         for i in range(2):
             for j in range(2):
-                outcomponents[idx] = wptype(2) * v2result[i, j]
+                outcomponents[idx] = wptype(2) * mat2[i, j]
                 idx = idx + 1
 
         for i in range(3):
             for j in range(3):
-                outcomponents[idx] = wptype(2) * v3result[i, j]
+                outcomponents[idx] = wptype(2) * mat3[i, j]
                 idx = idx + 1
 
         for i in range(4):
             for j in range(4):
-                outcomponents[idx] = wptype(2) * v4result[i, j]
+                outcomponents[idx] = wptype(2) * mat4[i, j]
                 idx = idx + 1
 
         for i in range(5):
             for j in range(5):
-                outcomponents[idx] = wptype(2) * v5result[i, j]
+                outcomponents[idx] = wptype(2) * mat5[i, j]
                 idx = idx + 1
 
-    kernel = getkernel(check_mat_add, suffix=dtype.__name__)
+    kernel = getkernel(check_mat_negation, suffix=dtype.__name__)
 
     if register_kernels:
         return
 
-    s2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    s3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    s4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    s5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    v2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    v3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    v4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    v5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
+    m2 = wp.array(randvals(rng, [1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
+    m3 = wp.array(randvals(rng, [1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
+    m4 = wp.array(randvals(rng, [1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
+    m5 = wp.array(randvals(rng, [1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
     outcomponents = wp.zeros(2 * 2 + 3 * 3 + 4 * 4 + 5 * 5, dtype=wptype, requires_grad=True, device=device)
 
-    wp.launch(
-        kernel,
-        dim=1,
-        inputs=[
-            s2,
-            s3,
-            s4,
-            s5,
-            v2,
-            v3,
-            v4,
-            v5,
-        ],
-        outputs=[outcomponents],
-        device=device,
-    )
+    wp.launch(kernel, dim=1, inputs=[m2, m3, m4, m5], outputs=[outcomponents], device=device)
 
-    assert_np_equal(outcomponents.numpy()[:4], 2 * (v2.numpy() + s2.numpy()).reshape(-1), tol=tol)
-    assert_np_equal(outcomponents.numpy()[4:13], 2 * (v3.numpy() + s3.numpy()).reshape(-1), tol=tol)
-    assert_np_equal(outcomponents.numpy()[13:29], 2 * (v4.numpy() + s4.numpy()).reshape(-1), tol=tol)
-    assert_np_equal(outcomponents.numpy()[29:54], 2 * (v5.numpy() + s5.numpy()).reshape(-1), tol=tol)
+    assert_np_equal(outcomponents.numpy()[:4], -2 * m2.numpy().reshape(-1), tol=tol)
+    assert_np_equal(outcomponents.numpy()[4:13], -2 * m3.numpy().reshape(-1), tol=tol)
+    assert_np_equal(outcomponents.numpy()[13:29], -2 * m4.numpy().reshape(-1), tol=tol)
+    assert_np_equal(outcomponents.numpy()[29:54], -2 * m5.numpy().reshape(-1), tol=tol)
 
     if dtype in np_float_types:
         idx = 0
         out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-        for dim, in1, in2 in [(2, s2, v2), (3, s3, v3), (4, s4, v4), (5, s5, v5)]:
+        for dim, input in [(2, m2), (3, m3), (4, m4), (5, m5)]:
             for i in range(dim):
                 for j in range(dim):
                     tape = wp.Tape()
                     with tape:
-                        wp.launch(
-                            kernel,
-                            dim=1,
-                            inputs=[
-                                s2,
-                                s3,
-                                s4,
-                                s5,
-                                v2,
-                                v3,
-                                v4,
-                                v5,
-                            ],
-                            outputs=[outcomponents],
-                            device=device,
-                        )
+                        wp.launch(kernel, dim=1, inputs=[m2, m3, m4, m5], outputs=[outcomponents], device=device)
                         wp.launch(
                             output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device
                         )
                     tape.backward(loss=out)
                     expectedresult = np.zeros((dim, dim), dtype=dtype)
-                    expectedresult[i, j] = 2
-                    assert_np_equal(tape.gradients[in2].numpy()[0], expectedresult, tol=10 * tol)
-                    expectedresult[i, j] = 2
-                    assert_np_equal(tape.gradients[in1].numpy()[0], expectedresult, tol=10 * tol)
+                    expectedresult[i, j] = -2
+                    assert_np_equal(tape.gradients[input].numpy()[0], expectedresult)
                     tape.zero()
-
                     idx = idx + 1
 
 
 def test_subtraction(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 5.0e-3,
@@ -2240,14 +517,14 @@ def test_subtraction(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    s2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    s3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    s4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    s5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    v2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    v3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    v4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    v5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
+    s2 = wp.array(randvals(rng, [1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
+    s3 = wp.array(randvals(rng, [1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
+    s4 = wp.array(randvals(rng, [1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
+    s5 = wp.array(randvals(rng, [1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
+    v2 = wp.array(randvals(rng, [1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
+    v3 = wp.array(randvals(rng, [1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
+    v4 = wp.array(randvals(rng, [1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
+    v5 = wp.array(randvals(rng, [1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
     outcomponents = wp.zeros(2 * 2 + 3 * 3 + 4 * 4 + 5 * 5, dtype=wptype, requires_grad=True, device=device)
 
     wp.launch(
@@ -2310,131 +587,8 @@ def test_subtraction(test, device, dtype, register_kernels=False):
                     idx = idx + 1
 
 
-def test_ddot(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 5.0e-3,
-        np.float32: 1.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-
-    def check_mat_dot(
-        s2: wp.array(dtype=mat22),
-        s3: wp.array(dtype=mat33),
-        s4: wp.array(dtype=mat44),
-        s5: wp.array(dtype=mat55),
-        v2: wp.array(dtype=mat22),
-        v3: wp.array(dtype=mat33),
-        v4: wp.array(dtype=mat44),
-        v5: wp.array(dtype=mat55),
-        dot2: wp.array(dtype=wptype),
-        dot3: wp.array(dtype=wptype),
-        dot4: wp.array(dtype=wptype),
-        dot5: wp.array(dtype=wptype),
-    ):
-        # multiply outputs by 2 so we've got something to backpropagate:
-        dot2[0] = wptype(2) * wp.ddot(v2[0], s2[0])
-        dot3[0] = wptype(2) * wp.ddot(v3[0], s3[0])
-        dot4[0] = wptype(2) * wp.ddot(v4[0], s4[0])
-        dot5[0] = wptype(2) * wp.ddot(v5[0], s5[0])
-
-    kernel = getkernel(check_mat_dot, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    s2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    s3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    s4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    s5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    v2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    v3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    v4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    v5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    dot2 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-    dot3 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-    dot4 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-    dot5 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-
-    tape = wp.Tape()
-    with tape:
-        wp.launch(
-            kernel,
-            dim=1,
-            inputs=[
-                s2,
-                s3,
-                s4,
-                s5,
-                v2,
-                v3,
-                v4,
-                v5,
-            ],
-            outputs=[dot2, dot3, dot4, dot5],
-            device=device,
-        )
-
-    assert_np_equal(dot2.numpy()[0], 2 * (v2.numpy() * s2.numpy()).sum(), tol=10 * tol)
-    assert_np_equal(dot3.numpy()[0], 2 * (v3.numpy() * s3.numpy()).sum(), tol=10 * tol)
-    assert_np_equal(dot4.numpy()[0], 2 * (v4.numpy() * s4.numpy()).sum(), tol=50 * tol)
-    assert_np_equal(dot5.numpy()[0], 2 * (v5.numpy() * s5.numpy()).sum(), tol=200 * tol)
-
-    if dtype in np_float_types:
-        tape.backward(loss=dot2)
-        sgrads = tape.gradients[s2].numpy()[0]
-        expected_grads = 2.0 * v2.numpy()[0]
-        assert_np_equal(sgrads, expected_grads, tol=10 * tol)
-
-        vgrads = tape.gradients[v2].numpy()[0]
-        expected_grads = 2.0 * s2.numpy()[0]
-        assert_np_equal(vgrads, expected_grads, tol=10 * tol)
-
-        tape.zero()
-
-        tape.backward(loss=dot3)
-        sgrads = tape.gradients[s3].numpy()[0]
-        expected_grads = 2.0 * v3.numpy()[0]
-        assert_np_equal(sgrads, expected_grads, tol=10 * tol)
-
-        vgrads = tape.gradients[v3].numpy()[0]
-        expected_grads = 2.0 * s3.numpy()[0]
-        assert_np_equal(vgrads, expected_grads, tol=10 * tol)
-
-        tape.zero()
-
-        tape.backward(loss=dot4)
-        sgrads = tape.gradients[s4].numpy()[0]
-        expected_grads = 2.0 * v4.numpy()[0]
-        assert_np_equal(sgrads, expected_grads, tol=10 * tol)
-
-        vgrads = tape.gradients[v4].numpy()[0]
-        expected_grads = 2.0 * s4.numpy()[0]
-        assert_np_equal(vgrads, expected_grads, tol=10 * tol)
-
-        tape.zero()
-
-        tape.backward(loss=dot5)
-        sgrads = tape.gradients[s5].numpy()[0]
-        expected_grads = 2.0 * v5.numpy()[0]
-        assert_np_equal(sgrads, expected_grads, tol=10 * tol)
-
-        vgrads = tape.gradients[v5].numpy()[0]
-        expected_grads = 2.0 * s5.numpy()[0]
-        assert_np_equal(vgrads, expected_grads, tol=10 * tol)
-
-        tape.zero()
-
-
 def test_determinant(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 5.0e-3,
@@ -2464,9 +618,9 @@ def test_determinant(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    v2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    v3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    v4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
+    v2 = wp.array(randvals(rng, [1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
+    v3 = wp.array(randvals(rng, [1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
+    v4 = wp.array(randvals(rng, [1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
     det2 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
     det3 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
     det4 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
@@ -2586,210 +740,114 @@ def test_determinant(test, device, dtype, register_kernels=False):
                     ],
                     outputs=[
                         det2,
-                        det3,
-                        det4,
-                    ],
-                    device=device,
-                )
-                dminus = det3.numpy()[0]
-                assert_np_equal((dplus - dminus) / (2.0 * dx * dplus), v3grads[i, j] / dplus, tol=fdtol)
-
-        for i in range(4):
-            for j in range(4):
-                v4test = v4.numpy()
-                v4test[0, i, j] += dx
-                wp.launch(
-                    kernel,
-                    dim=1,
-                    inputs=[
-                        v2,
-                        v3,
-                        wp.array(v4test, dtype=v4.dtype, requires_grad=True, device=device),
-                    ],
-                    outputs=[
-                        det2,
-                        det3,
-                        det4,
-                    ],
-                    device=device,
-                )
-                dplus = det4.numpy()[0]
-                v4test[0, i, j] -= 2.0 * dx
-                wp.launch(
-                    kernel,
-                    dim=1,
-                    inputs=[
-                        v2,
-                        v3,
-                        wp.array(v4test, dtype=v4.dtype, requires_grad=True, device=device),
-                    ],
-                    outputs=[
-                        det2,
-                        det3,
-                        det4,
-                    ],
-                    device=device,
-                )
-                dminus = det4.numpy()[0]
-                assert_np_equal((dplus - dminus) / (2.0 * dx * dplus), v4grads[i, j] / dplus, tol=fdtol)
-
-
-def test_trace(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 1.0e-3,
-        np.float32: 1.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-
-    def check_mat_trace(
-        v2: wp.array(dtype=mat22),
-        v3: wp.array(dtype=mat33),
-        v4: wp.array(dtype=mat44),
-        v5: wp.array(dtype=mat55),
-        tr2: wp.array(dtype=wptype),
-        tr3: wp.array(dtype=wptype),
-        tr4: wp.array(dtype=wptype),
-        tr5: wp.array(dtype=wptype),
-    ):
-        # multiply outputs by 2 so we've got something to backpropagate:
-        tr2[0] = wptype(2) * wp.trace(v2[0])
-        tr3[0] = wptype(2) * wp.trace(v3[0])
-        tr4[0] = wptype(2) * wp.trace(v4[0])
-        tr5[0] = wptype(2) * wp.trace(v5[0])
-
-    kernel = getkernel(check_mat_trace, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    v2 = wp.array(randvals([1, 2, 2], dtype), dtype=mat22, requires_grad=True, device=device)
-    v3 = wp.array(randvals([1, 3, 3], dtype), dtype=mat33, requires_grad=True, device=device)
-    v4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
-    v5 = wp.array(randvals([1, 5, 5], dtype), dtype=mat55, requires_grad=True, device=device)
-    tr2 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-    tr3 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-    tr4 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-    tr5 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-
-    tape = wp.Tape()
-    with tape:
-        wp.launch(
-            kernel,
-            dim=1,
-            inputs=[
-                v2,
-                v3,
-                v4,
-                v5,
-            ],
-            outputs=[
-                tr2,
-                tr3,
-                tr4,
-                tr5,
-            ],
-            device=device,
-        )
-
-    assert_np_equal(tr2.numpy()[0], 2 * np.trace(v2.numpy()[0]), tol=10 * tol)
-    assert_np_equal(tr3.numpy()[0], 2 * np.trace(v3.numpy()[0]), tol=10 * tol)
-    assert_np_equal(tr4.numpy()[0], 2 * np.trace(v4.numpy()[0]), tol=200 * tol)
-    assert_np_equal(tr4.numpy()[0], 2 * np.trace(v4.numpy()[0]), tol=200 * tol)
-
-    if dtype in np_float_types:
-        tape.backward(loss=tr2)
-        vgrads = tape.gradients[v2].numpy()[0]
-        assert_np_equal(vgrads, 2.0 * np.eye(2), tol=10 * tol)
-        tape.zero()
-
-        tape.backward(loss=tr3)
-        vgrads = tape.gradients[v3].numpy()[0]
-        assert_np_equal(vgrads, 2.0 * np.eye(3), tol=10 * tol)
-        tape.zero()
-
-        tape.backward(loss=tr4)
-        vgrads = tape.gradients[v4].numpy()[0]
-        assert_np_equal(vgrads, 2.0 * np.eye(4), tol=10 * tol)
-        tape.zero()
-
-        tape.backward(loss=tr5)
-        vgrads = tape.gradients[v5].numpy()[0]
-        assert_np_equal(vgrads, 2.0 * np.eye(5), tol=10 * tol)
-        tape.zero()
-
-
-def test_diag(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 1.0e-3,
-        np.float32: 1.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    vec5 = wp.types.vector(length=5, dtype=wptype)
-
-    output_select_kernel = get_select_kernel(wptype)
-
-    def check_mat_diag(
-        s5: wp.array(dtype=vec5),
-        outcomponents: wp.array(dtype=wptype),
-    ):
-        # multiply outputs by 2 so we've got something to backpropagate:
-        m55result = wptype(2) * wp.diag(s5[0])
-
-        idx = 0
-        for i in range(5):
-            for j in range(5):
-                outcomponents[idx] = m55result[i, j]
-                idx = idx + 1
-
-    kernel = getkernel(check_mat_diag, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    s5 = wp.array(randvals([1, 5], dtype), dtype=vec5, requires_grad=True, device=device)
-    outcomponents = wp.zeros(5 * 5, dtype=wptype, requires_grad=True, device=device)
-    out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-
-    wp.launch(kernel, dim=1, inputs=[s5], outputs=[outcomponents], device=device)
+                        det3,
+                        det4,
+                    ],
+                    device=device,
+                )
+                dminus = det3.numpy()[0]
+                assert_np_equal((dplus - dminus) / (2.0 * dx * dplus), v3grads[i, j] / dplus, tol=fdtol)
 
-    assert_np_equal(outcomponents.numpy(), 2 * np.diag(s5.numpy()[0]), tol=tol)
+        for i in range(4):
+            for j in range(4):
+                v4test = v4.numpy()
+                v4test[0, i, j] += dx
+                wp.launch(
+                    kernel,
+                    dim=1,
+                    inputs=[
+                        v2,
+                        v3,
+                        wp.array(v4test, dtype=v4.dtype, requires_grad=True, device=device),
+                    ],
+                    outputs=[
+                        det2,
+                        det3,
+                        det4,
+                    ],
+                    device=device,
+                )
+                dplus = det4.numpy()[0]
+                v4test[0, i, j] -= 2.0 * dx
+                wp.launch(
+                    kernel,
+                    dim=1,
+                    inputs=[
+                        v2,
+                        v3,
+                        wp.array(v4test, dtype=v4.dtype, requires_grad=True, device=device),
+                    ],
+                    outputs=[
+                        det2,
+                        det3,
+                        det4,
+                    ],
+                    device=device,
+                )
+                dminus = det4.numpy()[0]
+                assert_np_equal((dplus - dminus) / (2.0 * dx * dplus), v4grads[i, j] / dplus, tol=fdtol)
 
-    if dtype in np_float_types:
-        idx = 0
-        for i in range(5):
-            for j in range(5):
-                tape = wp.Tape()
-                with tape:
-                    wp.launch(kernel, dim=1, inputs=[s5], outputs=[outcomponents], device=device)
-                    wp.launch(output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device)
-                tape.backward(loss=out)
-                expectedresult = np.zeros(5, dtype=dtype)
-                if i == j:
-                    expectedresult[i] = 2
-                assert_np_equal(tape.gradients[s5].numpy()[0], expectedresult, tol=10 * tol)
-                tape.zero()
 
-                idx = idx + 1
+# Unused. Why?
+# def test_get_diag(test, device, dtype, register_kernels=False):
+#     tol = {
+#         np.float16: 1.0e-3,
+#         np.float32: 1.0e-6,
+#         np.float64: 1.0e-8,
+#     }.get(dtype, 0)
+#
+#     wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+#     mat55 = wp.types.vector(shape=(5, 5), dtype=wptype)
+#
+#     output_select_kernel = get_select_kernel(wptype)
+#
+#     def check_mat_diag(
+#         m55: wp.array(dtype=mat55),
+#         outcomponents: wp.array(dtype=wptype),
+#     ):
+#         # multiply outputs by 2 so we've got something to backpropagate:
+#         vec5result = wptype(2) * wp.get_diag(m55[0])
+#
+#         idx = 0
+#         for i in range(5):
+#             outcomponents[idx] = vec5result[i]
+#             idx = idx + 1
+#
+#     kernel = getkernel(check_mat_diag, suffix=dtype.__name__)
+#
+#     if register_kernels:
+#         return
+#
+#     m55 = wp.array(randvals((1, 5, 5), dtype), dtype=mat55, requires_grad=True, device=device)
+#     outcomponents = wp.zeros(5, dtype=wptype, requires_grad=True, device=device)
+#     out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
+#
+#     wp.launch(kernel, dim=1, inputs=[m55], outputs=[outcomponents], device=device)
+#
+#     assert_np_equal(outcomponents.numpy(), 2 * np.diag(m55.numpy()[0]), tol=tol)
+#
+#     if dtype in np_float_types:
+#         idx = 0
+#         for i in range(5):
+#             tape = wp.Tape()
+#             with tape:
+#                 wp.launch(kernel, dim=1, inputs=[m55], outputs=[outcomponents], device=device)
+#                 wp.launch(output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device)
+#             tape.backward(loss=out)
+#             expectedresult = np.zeros((5, 5), dtype=dtype)
+#             expectedresult[i, i] = 2
+#             assert_np_equal(tape.gradients[m55].numpy()[0], expectedresult, tol=10 * tol)
+#             tape.zero()
+#
+#             idx = idx + 1
 
 
 def test_inverse(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
-        np.float16: 2.0e-3,
-        np.float32: 1.0e-6,
+        np.float16: 5.0e-2,
+        np.float32: 1.0e-5,
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
@@ -2832,9 +890,15 @@ def test_inverse(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    m2 = wp.array(2 * (randvals([1, 2, 2], dtype) + 0.2 * np.eye(2)), dtype=mat22, requires_grad=True, device=device)
-    m3 = wp.array(2 * (randvals([1, 3, 3], dtype) + 0.2 * np.eye(3)), dtype=mat33, requires_grad=True, device=device)
-    m4 = wp.array(2 * (randvals([1, 4, 4], dtype) + 0.2 * np.eye(4)), dtype=mat44, requires_grad=True, device=device)
+    m2 = wp.array(
+        2 * (randvals(rng, [1, 2, 2], dtype) + 0.2 * np.eye(2)), dtype=mat22, requires_grad=True, device=device
+    )
+    m3 = wp.array(
+        2 * (randvals(rng, [1, 3, 3], dtype) + 0.2 * np.eye(3)), dtype=mat33, requires_grad=True, device=device
+    )
+    m4 = wp.array(
+        2 * (randvals(rng, [1, 4, 4], dtype) + 0.2 * np.eye(4)), dtype=mat44, requires_grad=True, device=device
+    )
 
     outcomponents = wp.zeros(2 * 2 + 3 * 3 + 4 * 4, dtype=wptype, requires_grad=True, device=device)
     out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
@@ -2949,7 +1013,7 @@ def test_inverse(test, device, dtype, register_kernels=False):
 
 
 def test_svd(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 1.0e-3,
@@ -3001,7 +1065,7 @@ def test_svd(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    m3 = wp.array(randvals([1, 3, 3], dtype) + np.eye(3), dtype=mat33, requires_grad=True, device=device)
+    m3 = wp.array(randvals(rng, [1, 3, 3], dtype) + np.eye(3), dtype=mat33, requires_grad=True, device=device)
 
     outcomponents = wp.zeros(2 * 3 * 3 + 3, dtype=wptype, requires_grad=True, device=device)
     Uout = wp.zeros(1, dtype=mat33, requires_grad=True, device=device)
@@ -3068,7 +1132,7 @@ def test_svd(test, device, dtype, register_kernels=False):
 
 
 def test_qr(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 2.0e-3,
@@ -3111,7 +1175,7 @@ def test_qr(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    m3 = wp.array(0.5 * (randvals([1, 3, 3], dtype) + np.eye(3)), dtype=mat33, requires_grad=True, device=device)
+    m3 = wp.array(0.5 * (randvals(rng, [1, 3, 3], dtype) + np.eye(3)), dtype=mat33, requires_grad=True, device=device)
 
     outcomponents = wp.zeros(2 * 3 * 3, dtype=wptype, requires_grad=True, device=device)
     Qout = wp.zeros(1, dtype=mat33, requires_grad=True, device=device)
@@ -3180,7 +1244,7 @@ def test_qr(test, device, dtype, register_kernels=False):
 
 
 def test_eig(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 4.0e-2,
@@ -3223,7 +1287,7 @@ def test_eig(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    m3_np = randvals([1, 3, 3], dtype) + np.eye(3, dtype=dtype)
+    m3_np = randvals(rng, [1, 3, 3], dtype) + np.eye(3, dtype=dtype)
     m3 = wp.array(m3_np, dtype=mat33, requires_grad=True, device=device)
 
     outcomponents = wp.zeros(3 * 3 + 3, dtype=wptype, requires_grad=True, device=device)
@@ -3292,7 +1356,7 @@ def test_eig(test, device, dtype, register_kernels=False):
 
 
 def test_skew(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 1.0e-3,
@@ -3323,7 +1387,7 @@ def test_skew(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    v3 = wp.array(randvals([1, 3], dtype), dtype=vec3, requires_grad=True, device=device)
+    v3 = wp.array(randvals(rng, [1, 3], dtype), dtype=vec3, requires_grad=True, device=device)
 
     outcomponents = wp.zeros(3 * 3, dtype=wptype, requires_grad=True, device=device)
     out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
@@ -3394,7 +1458,7 @@ def test_skew(test, device, dtype, register_kernels=False):
 
 
 def test_transform_point(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 5.0e-3,
@@ -3425,8 +1489,8 @@ def test_transform_point(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    v3 = wp.array(randvals([1, 3], dtype), dtype=vec3, requires_grad=True, device=device)
-    m4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
+    v3 = wp.array(randvals(rng, [1, 3], dtype), dtype=vec3, requires_grad=True, device=device)
+    m4 = wp.array(randvals(rng, [1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
 
     outcomponents = wp.zeros(3, dtype=wptype, requires_grad=True, device=device)
     out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
@@ -3455,7 +1519,7 @@ def test_transform_point(test, device, dtype, register_kernels=False):
 
 
 def test_transform_vector(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 5.0e-3,
@@ -3486,8 +1550,8 @@ def test_transform_vector(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    v3 = wp.array(randvals([1, 3], dtype), dtype=vec3, requires_grad=True, device=device)
-    m4 = wp.array(randvals([1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
+    v3 = wp.array(randvals(rng, [1, 3], dtype), dtype=vec3, requires_grad=True, device=device)
+    m4 = wp.array(randvals(rng, [1, 4, 4], dtype), dtype=mat44, requires_grad=True, device=device)
 
     outcomponents = wp.zeros(3, dtype=wptype, requires_grad=True, device=device)
     out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
@@ -3514,338 +1578,6 @@ def test_transform_vector(test, device, dtype, register_kernels=False):
             tape.zero()
 
 
-def test_anon_type_instance(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 5.0e-3,
-        np.float32: 1.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-
-    def check_scalar_init(
-        input: wp.array(dtype=wptype),
-        output: wp.array(dtype=wptype),
-    ):
-        m2result = wp.matrix(input[0], shape=(2, 2))
-        m3result = wp.matrix(input[1], shape=(3, 3))
-        m4result = wp.matrix(input[2], shape=(4, 4))
-        m5result = wp.matrix(input[3], shape=(5, 5))
-        m32result = wp.matrix(input[4], shape=(3, 2))
-
-        idx = 0
-        for i in range(2):
-            for j in range(2):
-                output[idx] = wptype(2) * m2result[i, j]
-                idx = idx + 1
-        for i in range(3):
-            for j in range(3):
-                output[idx] = wptype(2) * m3result[i, j]
-                idx = idx + 1
-        for i in range(4):
-            for j in range(4):
-                output[idx] = wptype(2) * m4result[i, j]
-                idx = idx + 1
-        for i in range(5):
-            for j in range(5):
-                output[idx] = wptype(2) * m5result[i, j]
-                idx = idx + 1
-        for i in range(3):
-            for j in range(2):
-                output[idx] = wptype(2) * m32result[i, j]
-                idx = idx + 1
-
-    def check_component_init(
-        input: wp.array(dtype=wptype),
-        output: wp.array(dtype=wptype),
-    ):
-        m2result = wp.matrix(input[0], input[1], input[2], input[3], shape=(2, 2))
-        m3result = wp.matrix(
-            input[4], input[5], input[6], input[7], input[8], input[9], input[10], input[11], input[12], shape=(3, 3)
-        )
-        m4result = wp.matrix(
-            input[13],
-            input[14],
-            input[15],
-            input[16],
-            input[17],
-            input[18],
-            input[19],
-            input[20],
-            input[21],
-            input[22],
-            input[23],
-            input[24],
-            input[25],
-            input[26],
-            input[27],
-            input[28],
-            shape=(4, 4),
-        )
-        m5result = wp.matrix(
-            input[29],
-            input[30],
-            input[31],
-            input[32],
-            input[33],
-            input[34],
-            input[35],
-            input[36],
-            input[37],
-            input[38],
-            input[39],
-            input[40],
-            input[41],
-            input[42],
-            input[43],
-            input[44],
-            input[45],
-            input[46],
-            input[47],
-            input[48],
-            input[49],
-            input[50],
-            input[51],
-            input[52],
-            input[53],
-            shape=(5, 5),
-        )
-        m32result = wp.matrix(input[54], input[55], input[56], input[57], input[58], input[59], shape=(3, 2))
-
-        idx = 0
-        for i in range(2):
-            for j in range(2):
-                output[idx] = wptype(2) * m2result[i, j]
-                idx = idx + 1
-        for i in range(3):
-            for j in range(3):
-                output[idx] = wptype(2) * m3result[i, j]
-                idx = idx + 1
-        for i in range(4):
-            for j in range(4):
-                output[idx] = wptype(2) * m4result[i, j]
-                idx = idx + 1
-        for i in range(5):
-            for j in range(5):
-                output[idx] = wptype(2) * m5result[i, j]
-                idx = idx + 1
-        for i in range(3):
-            for j in range(2):
-                output[idx] = wptype(2) * m32result[i, j]
-                idx = idx + 1
-
-    scalar_kernel = getkernel(check_scalar_init, suffix=dtype.__name__)
-    component_kernel = getkernel(check_component_init, suffix=dtype.__name__)
-    output_select_kernel = get_select_kernel(wptype)
-
-    if register_kernels:
-        return
-
-    input = wp.array(randvals([5], dtype), requires_grad=True, device=device)
-    output = wp.zeros(2 * 2 + 3 * 3 + 4 * 4 + 5 * 5 + 3 * 2, dtype=wptype, requires_grad=True, device=device)
-
-    wp.launch(scalar_kernel, dim=1, inputs=[input], outputs=[output], device=device)
-
-    assert_np_equal(output.numpy()[:4], 2 * np.array([input.numpy()[0]] * 2 * 2), tol=1.0e-6)
-    assert_np_equal(output.numpy()[4:13], 2 * np.array([input.numpy()[1]] * 3 * 3), tol=1.0e-6)
-    assert_np_equal(output.numpy()[13:29], 2 * np.array([input.numpy()[2]] * 4 * 4), tol=1.0e-6)
-    assert_np_equal(output.numpy()[29:54], 2 * np.array([input.numpy()[3]] * 5 * 5), tol=1.0e-6)
-    assert_np_equal(output.numpy()[54:], 2 * np.array([input.numpy()[4]] * 3 * 2), tol=1.0e-6)
-
-    if dtype in np_float_types:
-        out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-        for i in range(len(output)):
-            tape = wp.Tape()
-            with tape:
-                wp.launch(scalar_kernel, dim=1, inputs=[input], outputs=[output], device=device)
-                wp.launch(output_select_kernel, dim=1, inputs=[output, i], outputs=[out], device=device)
-
-            tape.backward(loss=out)
-            expected = np.zeros_like(input.numpy())
-            if i < 4:
-                expected[0] = 2
-            elif i < 13:
-                expected[1] = 2
-            elif i < 29:
-                expected[2] = 2
-            elif i < 54:
-                expected[3] = 2
-            else:
-                expected[4] = 2
-
-            assert_np_equal(tape.gradients[input].numpy(), expected, tol=tol)
-
-            tape.reset()
-            tape.zero()
-
-    input = wp.array(randvals([2 * 2 + 3 * 3 + 4 * 4 + 5 * 5 + 3 * 2], dtype), requires_grad=True, device=device)
-    output = wp.zeros(2 * 2 + 3 * 3 + 4 * 4 + 5 * 5 + 3 * 2, dtype=wptype, requires_grad=True, device=device)
-
-    wp.launch(component_kernel, dim=1, inputs=[input], outputs=[output], device=device)
-
-    assert_np_equal(output.numpy(), 2 * input.numpy(), tol=1.0e-6)
-
-    if dtype in np_float_types:
-        out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
-        for i in range(len(output)):
-            tape = wp.Tape()
-            with tape:
-                wp.launch(component_kernel, dim=1, inputs=[input], outputs=[output], device=device)
-                wp.launch(output_select_kernel, dim=1, inputs=[output, i], outputs=[out], device=device)
-
-            tape.backward(loss=out)
-            expected = np.zeros_like(input.numpy())
-            expected[i] = 2
-
-            assert_np_equal(tape.gradients[input].numpy(), expected, tol=tol)
-
-            tape.reset()
-            tape.zero()
-
-
-def test_identity(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
-
-    tol = {
-        np.float16: 5.0e-3,
-        np.float32: 1.0e-6,
-        np.float64: 1.0e-8,
-    }.get(dtype, 0)
-
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-
-    def check_identity_mat(
-        output: wp.array(dtype=wptype),
-    ):
-        m2result = wp.identity(dtype=wptype, n=2)
-        m3result = wp.identity(dtype=wptype, n=3)
-        m4result = wp.identity(dtype=wptype, n=4)
-        m5result = wp.identity(dtype=wptype, n=5)
-
-        idx = 0
-        for i in range(2):
-            for j in range(2):
-                output[idx] = wptype(2) * m2result[i, j]
-                idx = idx + 1
-        for i in range(3):
-            for j in range(3):
-                output[idx] = wptype(2) * m3result[i, j]
-                idx = idx + 1
-        for i in range(4):
-            for j in range(4):
-                output[idx] = wptype(2) * m4result[i, j]
-                idx = idx + 1
-        for i in range(5):
-            for j in range(5):
-                output[idx] = wptype(2) * m5result[i, j]
-                idx = idx + 1
-
-    id_kernel = getkernel(check_identity_mat, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    output = wp.zeros(2 * 2 + 3 * 3 + 4 * 4 + 5 * 5, dtype=wptype, requires_grad=True, device=device)
-    wp.launch(id_kernel, dim=1, inputs=[], outputs=[output], device=device)
-    assert_np_equal(output.numpy()[:4], 2 * np.eye(2), tol=1.0e-6)
-    assert_np_equal(output.numpy()[4:13], 2 * np.eye(3), tol=1.0e-6)
-    assert_np_equal(output.numpy()[13:29], 2 * np.eye(4), tol=1.0e-6)
-    assert_np_equal(output.numpy()[29:], 2 * np.eye(5), tol=1.0e-6)
-
-
-def test_equivalent_types(test, device, dtype, register_kernels=False):
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-
-    # matrix types
-    mat22 = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33 = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55 = wp.types.matrix(shape=(5, 5), dtype=wptype)
-
-    # matrix types equivalent to the above
-    mat22_equiv = wp.types.matrix(shape=(2, 2), dtype=wptype)
-    mat33_equiv = wp.types.matrix(shape=(3, 3), dtype=wptype)
-    mat44_equiv = wp.types.matrix(shape=(4, 4), dtype=wptype)
-    mat55_equiv = wp.types.matrix(shape=(5, 5), dtype=wptype)
-
-    # declare kernel with original types
-    def check_equivalence(
-        m2: mat22,
-        m3: mat33,
-        m4: mat44,
-        m5: mat55,
-    ):
-        wp.expect_eq(m2, mat22(wptype(42)))
-        wp.expect_eq(m3, mat33(wptype(43)))
-        wp.expect_eq(m4, mat44(wptype(44)))
-        wp.expect_eq(m5, mat55(wptype(45)))
-
-        wp.expect_eq(m2, mat22_equiv(wptype(42)))
-        wp.expect_eq(m3, mat33_equiv(wptype(43)))
-        wp.expect_eq(m4, mat44_equiv(wptype(44)))
-        wp.expect_eq(m5, mat55_equiv(wptype(45)))
-
-    kernel = getkernel(check_equivalence, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    # call kernel with equivalent types
-    m2 = mat22_equiv(42)
-    m3 = mat33_equiv(43)
-    m4 = mat44_equiv(44)
-    m5 = mat55_equiv(45)
-
-    wp.launch(kernel, dim=1, inputs=[m2, m3, m4, m5], device=device)
-
-
-def test_conversions(test, device, dtype, register_kernels=False):
-    def check_matrices_equal(
-        m0: wp.mat22,
-        m1: wp.mat22,
-        m2: wp.mat22,
-        m3: wp.mat22,
-        m4: wp.mat22,
-        m5: wp.mat22,
-        m6: wp.mat22,
-    ):
-        wp.expect_eq(m1, m0)
-        wp.expect_eq(m2, m0)
-        wp.expect_eq(m3, m0)
-        wp.expect_eq(m4, m0)
-        wp.expect_eq(m5, m0)
-        wp.expect_eq(m6, m0)
-
-    kernel = getkernel(check_matrices_equal, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    m0 = wp.mat22(1, 2, 3, 4)
-
-    # test explicit conversions - constructing matrices from different containers
-    m1 = wp.mat22(((1, 2), (3, 4)))  # nested tuples
-    m2 = wp.mat22([[1, 2], [3, 4]])  # nested lists
-    m3 = wp.mat22(np.array([[1, 2], [3, 4]], dtype=dtype))  # 2d array
-    m4 = wp.mat22((1, 2, 3, 4))  # flat tuple
-    m5 = wp.mat22([1, 2, 3, 4])  # flat list
-    m6 = wp.mat22(np.array([1, 2, 3, 4], dtype=dtype))  # 1d array
-
-    wp.launch(kernel, dim=1, inputs=[m0, m1, m2, m3, m4, m5, m6], device=device)
-
-    # test implicit conversions - passing different containers as matrices to wp.launch()
-    m1 = ((1, 2), (3, 4))  # nested tuples
-    m2 = [[1, 2], [3, 4]]  # nested lists
-    m3 = np.array([[1, 2], [3, 4]], dtype=dtype)  # 2d array
-    m4 = (1, 2, 3, 4)  # flat tuple
-    m5 = [1, 2, 3, 4]  # flat list
-    m6 = np.array([1, 2, 3, 4], dtype=dtype)  # 1d array
-
-    wp.launch(kernel, dim=1, inputs=[m0, m1, m2, m3, m4, m5, m6], device=device)
-
-
 # Test matrix constructors using explicit type (float16)
 # note that these tests are specifically not using generics / closure
 # args to create kernels dynamically (like the rest of this file)
@@ -3869,6 +1601,22 @@ def test_constructors_explicit_precision():
             wp.expect_eq(custom[i, j], wp.float16(i) * wp.float16(2.0) + wp.float16(j))
 
 
+mat32d = wp.mat(shape=(3, 2), dtype=wp.float64)
+
+
+@wp.kernel
+def test_matrix_constructor_value_func():
+    a = wp.mat22()
+    b = wp.matrix(a, shape=(2, 2))
+    c = mat32d()
+    d = mat32d(c, shape=(3, 2))
+    e = mat32d(wp.float64(1.0), wp.float64(2.0), wp.float64(1.0), wp.float64(2.0), wp.float64(1.0), wp.float64(2.0))
+    f = mat32d(
+        wp.vec3d(wp.float64(1.0), wp.float64(2.0), wp.float64(3.0)),
+        wp.vec3d(wp.float64(1.0), wp.float64(2.0), wp.float64(3.0)),
+    )
+
+
 # Same as above but with a default (float/int) type
 # which tests some different code paths that
 # need to ensure types are correctly canonicalized
@@ -3931,167 +1679,149 @@ def test_constructors_constant_shape():
             m[i, j] = float(i * j)
 
 
-def register(parent):
-    devices = get_test_devices()
-
-    class TestMat(parent):
-        pass
-
-    add_kernel_test(TestMat, test_constructors_explicit_precision, dim=1, devices=devices)
-    add_kernel_test(TestMat, test_constructors_default_precision, dim=1, devices=devices)
-    add_kernel_test(TestMat, test_constructors_constant_shape, dim=1, devices=devices)
-
-    mat103 = wp.types.matrix(shape=(10, 3), dtype=float)
-    add_kernel_test(
-        TestMat,
-        test_matrix_mutation,
-        dim=1,
-        inputs=[
-            mat103(
-                1.0,
-                2.0,
-                3.0,
-                2.0,
-                4.0,
-                6.0,
-                3.0,
-                6.0,
-                9.0,
-                4.0,
-                8.0,
-                12.0,
-                5.0,
-                10.0,
-                15.0,
-                6.0,
-                12.0,
-                18.0,
-                7.0,
-                14.0,
-                21.0,
-                8.0,
-                16.0,
-                24.0,
-                9.0,
-                18.0,
-                27.0,
-                10.0,
-                20.0,
-                30.0,
-            )
-        ],
-        devices=devices,
-    )
-
-    for dtype in np_signed_int_types + np_float_types:
-        add_function_test_register_kernel(
-            TestMat, f"test_negation_{dtype.__name__}", test_negation, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_subtraction_{dtype.__name__}", test_subtraction, devices=devices, dtype=dtype
-        )
-
-    for dtype in np_scalar_types:
-        add_function_test(TestMat, f"test_arrays_{dtype.__name__}", test_arrays, devices=devices, dtype=dtype)
-        add_function_test_register_kernel(
-            TestMat, f"test_constructors_{dtype.__name__}", test_constructors, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_anon_type_instance_{dtype.__name__}", test_anon_type_instance, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_identity_{dtype.__name__}", test_identity, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_indexing_{dtype.__name__}", test_indexing, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_equality_{dtype.__name__}", test_equality, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat,
-            f"test_scalar_multiplication_{dtype.__name__}",
-            test_scalar_multiplication,
-            devices=devices,
-            dtype=dtype,
-        )
-        add_function_test_register_kernel(
-            TestMat,
-            f"test_matvec_multiplication_{dtype.__name__}",
-            test_matvec_multiplication,
-            devices=devices,
-            dtype=dtype,
-        )
-        add_function_test_register_kernel(
-            TestMat,
-            f"test_matmat_multiplication_{dtype.__name__}",
-            test_matmat_multiplication,
-            devices=devices,
-            dtype=dtype,
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_cw_multiplication_{dtype.__name__}", test_cw_multiplication, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_cw_division_{dtype.__name__}", test_cw_division, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_outer_product_{dtype.__name__}", test_outer_product, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_transpose_{dtype.__name__}", test_transpose, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_scalar_division_{dtype.__name__}", test_scalar_division, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_addition_{dtype.__name__}", test_addition, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_ddot_{dtype.__name__}", test_ddot, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_trace_{dtype.__name__}", test_trace, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_diag_{dtype.__name__}", test_diag, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_equivalent_types_{dtype.__name__}", test_equivalent_types, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_conversions_{dtype.__name__}", test_conversions, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_constants_{dtype.__name__}", test_constants, devices=devices, dtype=dtype
+devices = get_test_devices()
+
+
+class TestMat(unittest.TestCase):
+    pass
+
+
+add_kernel_test(TestMat, test_constructors_explicit_precision, dim=1, devices=devices)
+add_kernel_test(TestMat, test_constructors_default_precision, dim=1, devices=devices)
+add_kernel_test(TestMat, test_constructors_constant_shape, dim=1, devices=devices)
+add_kernel_test(TestMat, test_matrix_constructor_value_func, dim=1, devices=devices)
+
+mat103 = wp.types.matrix(shape=(10, 3), dtype=float)
+add_kernel_test(
+    TestMat,
+    test_matrix_mutation,
+    dim=1,
+    inputs=[
+        mat103(
+            1.0,
+            2.0,
+            3.0,
+            2.0,
+            4.0,
+            6.0,
+            3.0,
+            6.0,
+            9.0,
+            4.0,
+            8.0,
+            12.0,
+            5.0,
+            10.0,
+            15.0,
+            6.0,
+            12.0,
+            18.0,
+            7.0,
+            14.0,
+            21.0,
+            8.0,
+            16.0,
+            24.0,
+            9.0,
+            18.0,
+            27.0,
+            10.0,
+            20.0,
+            30.0,
         )
+    ],
+    devices=devices,
+)
 
-    for dtype in np_float_types:
-        add_function_test_register_kernel(
-            TestMat, f"test_quat_constructor_{dtype.__name__}", test_quat_constructor, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_inverse_{dtype.__name__}", test_inverse, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(TestMat, f"test_svd_{dtype.__name__}", test_svd, devices=devices, dtype=dtype)
-        add_function_test_register_kernel(TestMat, f"test_qr_{dtype.__name__}", test_qr, devices=devices, dtype=dtype)
-        add_function_test_register_kernel(TestMat, f"test_eig_{dtype.__name__}", test_eig, devices=devices, dtype=dtype)
-        add_function_test_register_kernel(
-            TestMat, f"test_transform_point_{dtype.__name__}", test_transform_point, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_transform_vector_{dtype.__name__}", test_transform_vector, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_determinant_{dtype.__name__}", test_determinant, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestMat, f"test_skew_{dtype.__name__}", test_skew, devices=devices, dtype=dtype
-        )
+for dtype in np_signed_int_types + np_float_types:
+    add_function_test_register_kernel(
+        TestMat, f"test_negation_{dtype.__name__}", test_negation, devices=devices, dtype=dtype
+    )
+    add_function_test_register_kernel(
+        TestMat, f"test_subtraction_{dtype.__name__}", test_subtraction, devices=devices, dtype=dtype
+    )
 
-    return TestMat
+add_function_test(
+    TestMat,
+    "test_anon_constructor_error_shape_keyword_missing",
+    test_anon_constructor_error_shape_keyword_missing,
+    devices=devices,
+)
+add_function_test(
+    TestMat,
+    "test_anon_constructor_error_dtype_keyword_missing",
+    test_anon_constructor_error_dtype_keyword_missing,
+    devices=devices,
+)
+add_function_test(
+    TestMat,
+    "test_anon_constructor_error_shape_mismatch",
+    test_anon_constructor_error_shape_mismatch,
+    devices=devices,
+)
+add_function_test(
+    TestMat,
+    "test_anon_constructor_error_invalid_arg_count",
+    test_anon_constructor_error_invalid_arg_count,
+    devices=devices,
+)
+add_function_test(
+    TestMat,
+    "test_tpl_constructor_error_incompatible_sizes",
+    test_tpl_constructor_error_incompatible_sizes,
+    devices=devices,
+)
+add_function_test(
+    TestMat,
+    "test_tpl_constructor_error_invalid_scalar_type",
+    test_tpl_constructor_error_invalid_scalar_type,
+    devices=devices,
+)
+add_function_test(
+    TestMat,
+    "test_tpl_constructor_error_invalid_vector_count",
+    test_tpl_constructor_error_invalid_vector_count,
+    devices=devices,
+)
+add_function_test(
+    TestMat,
+    "test_tpl_constructor_error_invalid_vector_shape",
+    test_tpl_constructor_error_invalid_vector_shape,
+    devices=devices,
+)
+add_function_test(
+    TestMat,
+    "test_tpl_constructor_error_invalid_arg_count",
+    test_tpl_constructor_error_invalid_arg_count,
+    devices=devices,
+)
+add_function_test(TestMat, "test_tpl_ops_with_anon", test_tpl_ops_with_anon)
+
+for dtype in np_float_types:
+    add_function_test(
+        TestMat, f"test_py_arithmetic_ops_{dtype.__name__}", test_py_arithmetic_ops, devices=None, dtype=dtype
+    )
+    add_function_test_register_kernel(
+        TestMat, f"test_quat_constructor_{dtype.__name__}", test_quat_constructor, devices=devices, dtype=dtype
+    )
+    add_function_test_register_kernel(
+        TestMat, f"test_inverse_{dtype.__name__}", test_inverse, devices=devices, dtype=dtype
+    )
+    add_function_test_register_kernel(TestMat, f"test_svd_{dtype.__name__}", test_svd, devices=devices, dtype=dtype)
+    add_function_test_register_kernel(TestMat, f"test_qr_{dtype.__name__}", test_qr, devices=devices, dtype=dtype)
+    add_function_test_register_kernel(TestMat, f"test_eig_{dtype.__name__}", test_eig, devices=devices, dtype=dtype)
+    add_function_test_register_kernel(
+        TestMat, f"test_transform_point_{dtype.__name__}", test_transform_point, devices=devices, dtype=dtype
+    )
+    add_function_test_register_kernel(
+        TestMat, f"test_transform_vector_{dtype.__name__}", test_transform_vector, devices=devices, dtype=dtype
+    )
+    add_function_test_register_kernel(
+        TestMat, f"test_determinant_{dtype.__name__}", test_determinant, devices=devices, dtype=dtype
+    )
+    add_function_test_register_kernel(TestMat, f"test_skew_{dtype.__name__}", test_skew, devices=devices, dtype=dtype)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2, failfast=True)