warp-lang 1.2.2__py3-none-manylinux2014_aarch64.whl → 1.3.1__py3-none-manylinux2014_aarch64.whl

warp/tests/test_scalar_ops.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# 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.unittest_utils import *
+
+np_signed_int_types = [
+    np.int8,
+    np.int16,
+    np.int32,
+    np.int64,
+    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 test_py_arithmetic_ops(test, device, dtype):
+    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+
+    def make_scalar(value):
+        if wptype in wp.types.int_types:
+            # Cast to the correct integer type to simulate wrapping.
+            return wptype._type_(value).value
+
+        return value
+
+    a = wptype(1)
+    test.assertAlmostEqual(+a, make_scalar(1))
+    test.assertAlmostEqual(-a, make_scalar(-1))
+    test.assertAlmostEqual(a + wptype(5), make_scalar(6))
+    test.assertAlmostEqual(a - wptype(5), make_scalar(-4))
+
+    a = wptype(2)
+    test.assertAlmostEqual(a * wptype(2), make_scalar(4))
+    test.assertAlmostEqual(wptype(2) * a, make_scalar(4))
+    test.assertAlmostEqual(a / wptype(2), make_scalar(1))
+    test.assertAlmostEqual(wptype(24) / a, make_scalar(12))
+
+
+def test_py_math_ops(test, device, dtype):
+    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+
+    def make_scalar(value):
+        if wptype in wp.types.int_types:
+            # Cast to the correct integer type to simulate wrapping.
+            return wptype._type_(value).value
+
+        return value
+
+    a = wptype(1)
+    test.assertAlmostEqual(wp.abs(a), 1)
+
+    if dtype in np_float_types:
+        test.assertAlmostEqual(wp.sin(a), 0.84147098480789650488, places=3)
+        test.assertAlmostEqual(wp.radians(a), 0.01745329251994329577, places=5)
+
+
+devices = get_test_devices()
+
+
+class TestScalarOps(unittest.TestCase):
+    pass
+
+
+for dtype in np_scalar_types:
+    add_function_test(
+        TestScalarOps, f"test_py_arithmetic_ops_{dtype.__name__}", test_py_arithmetic_ops, devices=None, dtype=dtype
+    )
+    add_function_test(TestScalarOps, f"test_py_math_ops_{dtype.__name__}", test_py_math_ops, devices=None, dtype=dtype)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=True)

warp/tests/test_sim_grad.py

@@ -237,5 +237,5 @@ for int_type, int_name in enumerate(["featherstone", "semiimplicit"]):
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2, failfast=True)

warp/tests/test_sim_kinematics.py

@@ -85,5 +85,5 @@ add_function_test(TestSimKinematics, "test_fk_ik", test_fk_ik, devices=devices)
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2, failfast=True)

warp/tests/test_smoothstep.py

@@ -162,5 +162,5 @@ add_function_test(TestSmoothstep, "test_smoothstep", test_smoothstep, devices=de
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_sparse.py

@@ -158,6 +158,9 @@ def test_bsr_get_set_diag(test, device):
     diag = bsr_get_diag(diag_bsr)
     assert_np_equal(diag_scalar_np, diag.numpy(), tol=0.000001)
 
+    diag = bsr_get_diag(2.0 * diag_bsr)
+    assert_np_equal(2.0 * diag_scalar_np, diag.numpy(), tol=0.000001)
+
     # Uniform block diagonal
 
     with test.assertRaisesRegex(ValueError, "BsrMatrix block type must be either warp matrix or scalar"):
@@ -181,6 +184,58 @@ def test_bsr_get_set_diag(test, device):
     assert np.all(diag_csr.values.numpy() == np.ones(nrow, dtype=float))
 
 
+def test_bsr_split_merge(test, device):
+    rng = np.random.default_rng(123)
+
+    block_shape = (4, 2)
+    nrow = 4
+    ncol = 8
+    shape = (block_shape[0] * nrow, block_shape[1] * ncol)
+    n = 20
+
+    rows = wp.array(rng.integers(0, high=nrow, size=n, dtype=int), dtype=int, device=device)
+    cols = wp.array(rng.integers(0, high=ncol, size=n, dtype=int), dtype=int, device=device)
+    vals = wp.array(rng.random(size=(n, block_shape[0], block_shape[1])), dtype=float, device=device)
+
+    bsr = bsr_zeros(nrow, ncol, wp.types.matrix(shape=block_shape, dtype=float), device=device)
+    bsr_set_from_triplets(bsr, rows, cols, vals)
+    ref = _bsr_to_dense(bsr)
+
+    bsr_split = bsr_copy(bsr, block_shape=(2, 2))
+    test.assertEqual(bsr_split.block_size, 4)
+    res = _bsr_to_dense(bsr_split)
+    assert_np_equal(res, ref, 0.0001)
+
+    bsr_split = bsr_copy(bsr, block_shape=(1, 1))
+    test.assertEqual(bsr_split.block_size, 1)
+    res = _bsr_to_dense(bsr_split)
+    assert_np_equal(res, ref, 0.0001)
+
+    bsr_merge = bsr_copy(bsr, block_shape=(4, 4))
+    test.assertEqual(bsr_merge.block_size, 16)
+    res = _bsr_to_dense(bsr_merge)
+    assert_np_equal(res, ref, 0.0001)
+
+    bsr_merge = bsr_copy(bsr, block_shape=(8, 8))
+    test.assertEqual(bsr_merge.block_size, 64)
+    res = _bsr_to_dense(bsr_merge)
+    assert_np_equal(res, ref, 0.0001)
+
+    with test.assertRaisesRegex(ValueError, "Incompatible dest and src block shapes"):
+        bsr_copy(bsr, block_shape=(3, 3))
+
+    with test.assertRaisesRegex(
+        ValueError, r"Dest block shape \(5, 5\) is not an exact multiple of src block shape \(4, 2\)"
+    ):
+        bsr_copy(bsr, block_shape=(5, 5))
+
+    with test.assertRaisesRegex(
+        ValueError,
+        "The total rows and columns of the src matrix cannot be evenly divided using the requested block shape",
+    ):
+        bsr_copy(bsr, block_shape=(32, 32))
+
+
 def make_test_bsr_transpose(block_shape, scalar_type):
     def test_bsr_transpose(test, device):
         rng = np.random.default_rng(123)
@@ -197,14 +252,11 @@ def make_test_bsr_transpose(block_shape, scalar_type):
 
         bsr = bsr_zeros(nrow, ncol, wp.types.matrix(shape=block_shape, dtype=scalar_type), device=device)
         bsr_set_from_triplets(bsr, rows, cols, vals)
-        ref = np.transpose(_bsr_to_dense(bsr))
+        ref = 2.0 * np.transpose(_bsr_to_dense(bsr))
 
-        bsr_transposed = bsr_zeros(
-            ncol, nrow, wp.types.matrix(shape=block_shape[::-1], dtype=scalar_type), device=device
-        )
-        bsr_set_transpose(dest=bsr_transposed, src=bsr)
+        bsr_transposed = (2.0 * bsr).transpose()
 
-        res = _bsr_to_dense(bsr_transposed)
+        res = _bsr_to_dense(bsr_transposed.eval())
         assert_np_equal(res, ref, 0.0001)
 
         if block_shape[0] != block_shape[-1]:
@@ -245,17 +297,14 @@ def make_test_bsr_axpy(block_shape, scalar_type):
         work_arrays = bsr_axpy_work_arrays()
         for alpha, beta in zip(alphas, betas):
             ref = alpha * _bsr_to_dense(x) + beta * _bsr_to_dense(y)
-            if beta == 0.0:
-                y = bsr_axpy(x, alpha=alpha, beta=beta, work_arrays=work_arrays)
-            else:
-                bsr_axpy(x, y, alpha, beta, work_arrays=work_arrays)
+            bsr_axpy(x, y, alpha, beta, work_arrays=work_arrays)
 
             res = _bsr_to_dense(y)
             assert_np_equal(res, ref, 0.0001)
 
         # test aliasing
         ref = 3.0 * _bsr_to_dense(y)
-        bsr_axpy(y, y, alpha=1.0, beta=2.0)
+        y += y * 2.0
         res = _bsr_to_dense(y)
         assert_np_equal(res, ref, 0.0001)
 
@@ -285,7 +334,7 @@ def make_test_bsr_mm(block_shape, scalar_type):
 
         nnz = 6
 
-        alphas = [-1.0, 0.0, 1.0]
+        alphas = [-1.0, 0.0, 2.0]
         betas = [2.0, -1.0, 0.0]
 
         x_rows = wp.array(rng.integers(0, high=x_nrow, size=nnz, dtype=int), dtype=int, device=device)
@@ -321,6 +370,15 @@ def make_test_bsr_mm(block_shape, scalar_type):
             res = _bsr_to_dense(z)
             assert_np_equal(res, ref, 0.0001)
 
+        # test reusing topology from work arrays
+        # (assumes betas[-1] = 0)
+        bsr_mm(x, y, z, alpha, beta, work_arrays=work_arrays, reuse_topology=True)
+        assert_np_equal(res, ref, 0.0001)
+
+        # using overloaded operators
+        x = (alpha * x) @ y
+        assert_np_equal(res, ref, 0.0001)
+
         # test aliasing of matrix arguments
         # x = alpha * z * x + beta * x
         alpha, beta = alphas[0], betas[0]
@@ -389,16 +447,24 @@ def make_test_bsr_mv(block_shape, scalar_type):
         for alpha, beta in zip(alphas, betas):
             ref = alpha * _bsr_to_dense(A) @ x.numpy().flatten() + beta * y.numpy().flatten()
             if beta == 0.0:
-                y = bsr_mv(A, x, alpha=alpha, beta=beta, work_buffer=work_buffer)
+                y = A @ x
             else:
                 bsr_mv(A, x, y, alpha, beta, work_buffer=work_buffer)
 
             res = y.numpy().flatten()
             assert_np_equal(res, ref, 0.0001)
 
+        # test transposed product
+        ref = alpha * y.numpy().flatten() @ _bsr_to_dense(A)
+        x = y @ (A * alpha)
+        res = x.numpy().flatten()
+        assert_np_equal(res, ref, 0.0001)
+
         # test aliasing
-        alpha, beta = alphas[0], betas[0]
         AAt = bsr_mm(A, bsr_transposed(A))
+        assert_np_equal(_bsr_to_dense(AAt), _bsr_to_dense(A) @ _bsr_to_dense(A).T, 0.0001)
+
+        alpha, beta = alphas[0], betas[0]
         ref = alpha * _bsr_to_dense(AAt) @ y.numpy().flatten() + beta * y.numpy().flatten()
         bsr_mv(AAt, y, y, alpha, beta)
         res = y.numpy().flatten()
@@ -443,6 +509,7 @@ class TestSparse(unittest.TestCase):
 add_function_test(TestSparse, "test_csr_from_triplets", test_csr_from_triplets, devices=devices)
 add_function_test(TestSparse, "test_bsr_from_triplets", test_bsr_from_triplets, devices=devices)
 add_function_test(TestSparse, "test_bsr_get_diag", test_bsr_get_set_diag, devices=devices)
+add_function_test(TestSparse, "test_bsr_split_merge", test_bsr_split_merge, devices=devices)
 
 add_function_test(TestSparse, "test_csr_transpose", make_test_bsr_transpose((1, 1), wp.float32), devices=devices)
 add_function_test(TestSparse, "test_bsr_transpose_1_3", make_test_bsr_transpose((1, 3), wp.float32), devices=devices)
@@ -462,5 +529,5 @@ add_function_test(TestSparse, "test_bsr_mv_3_3", make_test_bsr_mv((3, 3), wp.flo
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_spatial.py

@@ -2142,5 +2142,5 @@ for dtype in np_float_types:
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_special_values.py

@@ -279,7 +279,6 @@ def test_is_special_quat(test, device, dtype, register_kernels=False):
 def test_is_special_int(test, device, dtype, register_kernels=False):
     vector_type = wp.types.vector(5, dtype)
     matrix_type = wp.types.matrix((5, 5), dtype)
-    quat_type = wp.types.quaternion(dtype)
 
     def check_is_special_int(bool_outputs: wp.array(dtype=wp.bool)):
         bool_outputs[0] = wp.isfinite(dtype(0))
@@ -294,16 +293,12 @@ def test_is_special_int(test, device, dtype, register_kernels=False):
         bool_outputs[7] = wp.isnan(matrix_type())
         bool_outputs[8] = wp.isinf(matrix_type())
 
-        bool_outputs[9] = wp.isfinite(quat_type())
-        bool_outputs[10] = wp.isnan(quat_type())
-        bool_outputs[11] = wp.isinf(quat_type())
-
     kernel = getkernel(check_is_special_int, suffix=dtype.__name__)
 
     if register_kernels:
         return
 
-    outputs_bool = wp.empty(12, dtype=wp.bool, device=device)
+    outputs_bool = wp.empty(9, dtype=wp.bool, device=device)
 
     wp.launch(kernel, dim=1, inputs=[outputs_bool], device=device)
 
@@ -321,10 +316,6 @@ def test_is_special_int(test, device, dtype, register_kernels=False):
     test.assertFalse(outputs_bool_cpu[7], "wp.isinf(matrix) is not False")
     test.assertFalse(outputs_bool_cpu[8], "wp.isnan(matrix) is not False")
 
-    test.assertTrue(outputs_bool_cpu[9], "wp.isfinite(quat) is not True")
-    test.assertFalse(outputs_bool_cpu[10], "wp.isinf(quat) is not False")
-    test.assertFalse(outputs_bool_cpu[11], "wp.isnan(quat) is not False")
-
 
 devices = get_test_devices()
 
@@ -358,5 +349,5 @@ for dtype in wp.types.int_types:
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2, failfast=False)

warp/tests/test_streams.py

@@ -464,6 +464,16 @@ class TestStreams(unittest.TestCase):
             # check results
             assert_np_equal(c0.numpy(), np.full(N, fill_value=2 * num_iters))
 
+    def test_stream_new_del(self):
+        # test the scenario in which a Stream is created but not initialized before gc
+        instance = wp.Stream.__new__(wp.Stream)
+        instance.__del__()
+
+    def test_event_new_del(self):
+        # test the scenario in which an Event is created but not initialized before gc
+        instance = wp.Event.__new__(wp.Event)
+        instance.__del__()
+
 
 add_function_test(TestStreams, "test_stream_set", test_stream_set, devices=devices)
 add_function_test(TestStreams, "test_stream_arg_explicit_sync", test_stream_arg_explicit_sync, devices=devices)
@@ -480,5 +490,5 @@ add_function_test(TestStreams, "test_event_synchronize", test_event_synchronize,
 add_function_test(TestStreams, "test_event_elapsed_time", test_event_elapsed_time, devices=devices)
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_struct.py

@@ -704,5 +704,5 @@ add_kernel_test(
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_tape.py

@@ -167,5 +167,5 @@ add_function_test(TestTape, "test_tape_visualize", test_tape_visualize, devices=
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_torch.py

@@ -25,6 +25,18 @@ def inc(a: wp.array(dtype=float)):
     a[tid] = a[tid] + 1.0
 
 
+@wp.kernel
+def inc_vector(a: wp.array(dtype=wp.vec3f)):
+    tid = wp.tid()
+    a[tid] = a[tid] + wp.vec3f(1.0)
+
+
+@wp.kernel
+def inc_matrix(a: wp.array(dtype=wp.mat22f)):
+    tid = wp.tid()
+    a[tid] = a[tid] + wp.mat22f(1.0)
+
+
 @wp.kernel
 def arange(start: int, step: int, a: wp.array(dtype=int)):
     tid = wp.tid()
@@ -216,6 +228,160 @@ def test_from_torch(test, device):
     wrap_mat_tensor_with_grad(6, 6, wp.spatial_matrix)
 
 
+def test_array_ctype_from_torch(test, device):
+    import torch
+
+    torch_device = wp.device_to_torch(device)
+
+    # automatically determine warp dtype
+    def wrap_scalar_tensor_implicit(torch_dtype):
+        t = torch.zeros(10, dtype=torch_dtype, device=torch_device)
+        a = wp.from_torch(t, return_ctype=True)
+        warp_dtype = wp.dtype_from_torch(torch_dtype)
+        ctype_size = ctypes.sizeof(warp_dtype._type_)
+        assert a.data == t.data_ptr()
+        assert a.grad == 0
+        assert a.ndim == 1
+        assert a.shape[0] == t.shape[0]
+        assert a.strides[0] == t.stride()[0] * ctype_size
+
+    wrap_scalar_tensor_implicit(torch.float64)
+    wrap_scalar_tensor_implicit(torch.float32)
+    wrap_scalar_tensor_implicit(torch.float16)
+    wrap_scalar_tensor_implicit(torch.int64)
+    wrap_scalar_tensor_implicit(torch.int32)
+    wrap_scalar_tensor_implicit(torch.int16)
+    wrap_scalar_tensor_implicit(torch.int8)
+    wrap_scalar_tensor_implicit(torch.uint8)
+    wrap_scalar_tensor_implicit(torch.bool)
+
+    # explicitly specify warp dtype
+    def wrap_scalar_tensor_explicit(torch_dtype, warp_dtype):
+        t = torch.zeros(10, dtype=torch_dtype, device=torch_device)
+        a = wp.from_torch(t, dtype=warp_dtype, return_ctype=True)
+        ctype_size = ctypes.sizeof(warp_dtype._type_)
+        assert a.data == t.data_ptr()
+        assert a.grad == 0
+        assert a.ndim == 1
+        assert a.shape[0] == t.shape[0]
+        assert a.strides[0] == t.stride()[0] * ctype_size
+
+    wrap_scalar_tensor_explicit(torch.float64, wp.float64)
+    wrap_scalar_tensor_explicit(torch.float32, wp.float32)
+    wrap_scalar_tensor_explicit(torch.float16, wp.float16)
+    wrap_scalar_tensor_explicit(torch.int64, wp.int64)
+    wrap_scalar_tensor_explicit(torch.int64, wp.uint64)
+    wrap_scalar_tensor_explicit(torch.int32, wp.int32)
+    wrap_scalar_tensor_explicit(torch.int32, wp.uint32)
+    wrap_scalar_tensor_explicit(torch.int16, wp.int16)
+    wrap_scalar_tensor_explicit(torch.int16, wp.uint16)
+    wrap_scalar_tensor_explicit(torch.int8, wp.int8)
+    wrap_scalar_tensor_explicit(torch.int8, wp.uint8)
+    wrap_scalar_tensor_explicit(torch.uint8, wp.uint8)
+    wrap_scalar_tensor_explicit(torch.uint8, wp.int8)
+    wrap_scalar_tensor_explicit(torch.bool, wp.uint8)
+    wrap_scalar_tensor_explicit(torch.bool, wp.int8)
+    wrap_scalar_tensor_explicit(torch.bool, wp.bool)
+
+    def wrap_vec_tensor(vec_dtype):
+        t = torch.zeros((10, vec_dtype._length_), dtype=torch.float32, device=torch_device)
+        a = wp.from_torch(t, dtype=vec_dtype, return_ctype=True)
+        ctype_size = ctypes.sizeof(vec_dtype._type_)
+        assert a.data == t.data_ptr()
+        assert a.grad == 0
+        assert a.ndim == 1
+        assert a.shape[0] == t.shape[0]
+        assert a.strides[0] == t.stride()[0] * ctype_size
+
+    wrap_vec_tensor(wp.vec2)
+    wrap_vec_tensor(wp.vec3)
+    wrap_vec_tensor(wp.vec4)
+    wrap_vec_tensor(wp.spatial_vector)
+    wrap_vec_tensor(wp.transform)
+
+    def wrap_mat_tensor(mat_dtype):
+        t = torch.zeros((10, *mat_dtype._shape_), dtype=torch.float32, device=torch_device)
+        a = wp.from_torch(t, dtype=mat_dtype, return_ctype=True)
+        ctype_size = ctypes.sizeof(mat_dtype._type_)
+        assert a.data == t.data_ptr()
+        assert a.grad == 0
+        assert a.ndim == 1
+        assert a.shape[0] == t.shape[0]
+        assert a.strides[0] == t.stride()[0] * ctype_size
+
+    wrap_mat_tensor(wp.mat22)
+    wrap_mat_tensor(wp.mat33)
+    wrap_mat_tensor(wp.mat44)
+    wrap_mat_tensor(wp.spatial_matrix)
+
+    def wrap_vec_tensor_with_existing_grad(vec_dtype):
+        t = torch.zeros((10, vec_dtype._length_), dtype=torch.float32, device=torch_device, requires_grad=True)
+        t.grad = torch.zeros((10, vec_dtype._length_), dtype=torch.float32, device=torch_device)
+        a = wp.from_torch(t, dtype=vec_dtype, return_ctype=True)
+        ctype_size = ctypes.sizeof(vec_dtype._type_)
+        assert a.data == t.data_ptr()
+        assert a.grad == t.grad.data_ptr()
+        assert a.ndim == 1
+        assert a.shape[0] == t.shape[0]
+        assert a.strides[0] == t.stride()[0] * ctype_size
+
+    wrap_vec_tensor_with_existing_grad(wp.vec2)
+    wrap_vec_tensor_with_existing_grad(wp.vec3)
+    wrap_vec_tensor_with_existing_grad(wp.vec4)
+    wrap_vec_tensor_with_existing_grad(wp.spatial_vector)
+    wrap_vec_tensor_with_existing_grad(wp.transform)
+
+    def wrap_vec_tensor_with_new_grad(vec_dtype):
+        t = torch.zeros((10, vec_dtype._length_), dtype=torch.float32, device=torch_device)
+        a = wp.from_torch(t, dtype=vec_dtype, requires_grad=True, return_ctype=True)
+        ctype_size = ctypes.sizeof(vec_dtype._type_)
+        assert a.data == t.data_ptr()
+        assert a.grad == t.grad.data_ptr()
+        assert a.ndim == 1
+        assert a.shape[0] == t.shape[0]
+        assert a.strides[0] == t.stride()[0] * ctype_size
+
+    wrap_vec_tensor_with_new_grad(wp.vec2)
+    wrap_vec_tensor_with_new_grad(wp.vec3)
+    wrap_vec_tensor_with_new_grad(wp.vec4)
+    wrap_vec_tensor_with_new_grad(wp.spatial_vector)
+    wrap_vec_tensor_with_new_grad(wp.transform)
+
+    def wrap_vec_tensor_with_torch_grad(vec_dtype):
+        t = torch.zeros((10, vec_dtype._length_), dtype=torch.float32, device=torch_device)
+        grad = torch.zeros((10, vec_dtype._length_), dtype=torch.float32, device=torch_device)
+        a = wp.from_torch(t, dtype=vec_dtype, grad=grad, return_ctype=True)
+        ctype_size = ctypes.sizeof(vec_dtype._type_)
+        assert a.data == t.data_ptr()
+        assert a.grad == grad.data_ptr()
+        assert a.ndim == 1
+        assert a.shape[0] == t.shape[0]
+        assert a.strides[0] == t.stride()[0] * ctype_size
+
+    wrap_vec_tensor_with_torch_grad(wp.vec2)
+    wrap_vec_tensor_with_torch_grad(wp.vec3)
+    wrap_vec_tensor_with_torch_grad(wp.vec4)
+    wrap_vec_tensor_with_torch_grad(wp.spatial_vector)
+    wrap_vec_tensor_with_torch_grad(wp.transform)
+
+    def wrap_vec_tensor_with_warp_grad(vec_dtype):
+        t = torch.zeros((10, vec_dtype._length_), dtype=torch.float32, device=torch_device)
+        grad = wp.zeros(10, dtype=vec_dtype, device=device)
+        a = wp.from_torch(t, dtype=vec_dtype, grad=grad, return_ctype=True)
+        ctype_size = ctypes.sizeof(vec_dtype._type_)
+        assert a.data == t.data_ptr()
+        assert a.grad == grad.ptr
+        assert a.ndim == 1
+        assert a.shape[0] == t.shape[0]
+        assert a.strides[0] == t.stride()[0] * ctype_size
+
+    wrap_vec_tensor_with_warp_grad(wp.vec2)
+    wrap_vec_tensor_with_warp_grad(wp.vec3)
+    wrap_vec_tensor_with_warp_grad(wp.vec4)
+    wrap_vec_tensor_with_warp_grad(wp.spatial_vector)
+    wrap_vec_tensor_with_warp_grad(wp.transform)
+
+
 def test_to_torch(test, device):
     import torch
 
@@ -659,6 +825,29 @@ def test_warp_graph_torch_stream(test, device):
     assert passed.item()
 
 
+def test_direct(test, device):
+    """Pass Torch tensors to Warp kernels directly"""
+
+    import torch
+
+    torch_device = wp.device_to_torch(device)
+    n = 12
+
+    s = torch.arange(n, dtype=torch.float32, device=torch_device)
+    v = torch.arange(n, dtype=torch.float32, device=torch_device).reshape((n // 3, 3))
+    m = torch.arange(n, dtype=torch.float32, device=torch_device).reshape((n // 4, 2, 2))
+
+    wp.launch(inc, dim=n, inputs=[s], device=device)
+    wp.launch(inc_vector, dim=n // 3, inputs=[v], device=device)
+    wp.launch(inc_matrix, dim=n // 4, inputs=[m], device=device)
+
+    expected = torch.arange(1, n + 1, dtype=torch.float32, device=torch_device)
+
+    assert torch.equal(s, expected)
+    assert torch.equal(v.reshape(n), expected)
+    assert torch.equal(m.reshape(n), expected)
+
+
 class TestTorch(unittest.TestCase):
     pass
 
@@ -690,6 +879,9 @@ try:
         add_function_test(TestTorch, "test_device_conversion", test_device_conversion, devices=torch_compatible_devices)
         add_function_test(TestTorch, "test_from_torch", test_from_torch, devices=torch_compatible_devices)
         add_function_test(TestTorch, "test_from_torch_slices", test_from_torch_slices, devices=torch_compatible_devices)
+        add_function_test(
+            TestTorch, "test_array_ctype_from_torch", test_array_ctype_from_torch, devices=torch_compatible_devices
+        )
         add_function_test(
             TestTorch,
             "test_from_torch_zero_strides",
@@ -699,6 +891,7 @@ try:
         add_function_test(TestTorch, "test_to_torch", test_to_torch, devices=torch_compatible_devices)
         add_function_test(TestTorch, "test_torch_zerocopy", test_torch_zerocopy, devices=torch_compatible_devices)
         add_function_test(TestTorch, "test_torch_autograd", test_torch_autograd, devices=torch_compatible_devices)
+        add_function_test(TestTorch, "test_direct", test_direct, devices=torch_compatible_devices)
 
     if torch_compatible_cuda_devices:
         add_function_test(
@@ -737,5 +930,5 @@ except Exception as e:
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_transient_module.py

@@ -81,5 +81,5 @@ class TestTransientModule(unittest.TestCase):
 add_function_test(TestTransientModule, "test_transient_module", test_transient_module, devices=devices)
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_types.py

@@ -550,5 +550,5 @@ for dtype in tuple(wp.types.scalar_types) + (int, float):
     add_function_test(TestTypes, f"test_vector_{dtype.__name__}", test_vector, devices=devices, dtype=dtype)
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_utils.py

@@ -484,5 +484,5 @@ add_function_test(
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2)