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

warp/tests/test_vec.py

@@ -49,23 +49,6 @@ def getkernel(func, suffix=""):
     return kernel_cache[key]
 
 
-def test_anon_constructor_error_dtype_keyword_missing(test, device):
-    @wp.kernel
-    def kernel():
-        wp.vector(length=123)
-
-    with test.assertRaisesRegex(
-        RuntimeError,
-        r"vec\(\) must have dtype as a keyword argument if it has no positional arguments, e.g.: wp.vector\(length=5, dtype=wp.float32\)$",
-    ):
-        wp.launch(
-            kernel,
-            dim=1,
-            inputs=[],
-            device=device,
-        )
-
-
 def test_anon_constructor_error_length_mismatch(test, device):
     @wp.kernel
     def kernel():
@@ -77,7 +60,7 @@ def test_anon_constructor_error_length_mismatch(test, device):
 
     with test.assertRaisesRegex(
         RuntimeError,
-        r"Incompatible vector lengths for casting copy constructor, 3 vs 2$",
+        r"incompatible vector of length 3 given when copy constructing a vector of length 2$",
     ):
         wp.launch(
             kernel,
@@ -87,14 +70,14 @@ def test_anon_constructor_error_length_mismatch(test, device):
         )
 
 
-def test_anon_constructor_error_numeric_arg_missing_1(test, device):
+def test_anon_constructor_error_numeric_arg_missing(test, device):
     @wp.kernel
     def kernel():
         wp.vector(1.0, 2.0, length=12345)
 
     with test.assertRaisesRegex(
         RuntimeError,
-        r"vec\(\) must have one scalar argument or the dtype keyword argument if the length keyword argument is specified, e.g.: wp.vec\(1.0, length=5\)$",
+        r"incompatible number of values given \(2\) when constructing a vector of length 12345$",
     ):
         wp.launch(
             kernel,
@@ -104,31 +87,14 @@ def test_anon_constructor_error_numeric_arg_missing_1(test, device):
         )
 
 
-def test_anon_constructor_error_numeric_arg_missing_2(test, device):
+def test_anon_constructor_error_length_arg_missing(test, device):
     @wp.kernel
     def kernel():
         wp.vector()
 
     with test.assertRaisesRegex(
         RuntimeError,
-        r"vec\(\) must have at least one numeric argument, if it's length, dtype is not specified$",
-    ):
-        wp.launch(
-            kernel,
-            dim=1,
-            inputs=[],
-            device=device,
-        )
-
-
-def test_anon_constructor_error_dtype_keyword_extraneous(test, device):
-    @wp.kernel
-    def kernel():
-        wp.vector(1.0, 2.0, 3.0, dtype=float)
-
-    with test.assertRaisesRegex(
-        RuntimeError,
-        r"vec\(\) should not have dtype specified if numeric arguments are given, the dtype will be inferred from the argument types$",
+        r"the `length` argument must be specified when zero-initializing a vector$",
     ):
         wp.launch(
             kernel,
@@ -145,9 +111,7 @@ def test_anon_constructor_error_numeric_args_mismatch(test, device):
 
     with test.assertRaisesRegex(
         RuntimeError,
-        r"All numeric arguments to vec\(\) constructor should have the same "
-        r"type, expected 2 arg_types of type <class 'warp.types.float32'>, "
-        r"received <class 'warp.types.float32'>,<class 'warp.types.int32'>$",
+        r"all values given when constructing a vector must have the same type$",
     ):
         wp.launch(
             kernel,
@@ -162,7 +126,9 @@ def test_tpl_constructor_error_incompatible_sizes(test, device):
     def kernel():
         wp.vec3(wp.vec2(1.0, 2.0))
 
-    with test.assertRaisesRegex(RuntimeError, r"Incompatible matrix sizes for casting copy constructor, 3 vs 2"):
+    with test.assertRaisesRegex(
+        RuntimeError, "incompatible vector of length 3 given when copy constructing a vector of length 2"
+    ):
         wp.launch(
             kernel,
             dim=1,
@@ -178,9 +144,7 @@ def test_tpl_constructor_error_numeric_args_mismatch(test, device):
 
     with test.assertRaisesRegex(
         RuntimeError,
-        r"All numeric arguments to vec\(\) constructor should have the same "
-        r"type, expected 2 arg_types of type <class 'warp.types.float32'>, "
-        r"received <class 'warp.types.float32'>,<class 'warp.types.int32'>$",
+        r"all values given when constructing a vector must have the same type$",
     ):
         wp.launch(
             kernel,
@@ -1207,12 +1171,6 @@ for dtype in np_float_types:
         dtype=dtype,
     )
 
-add_function_test(
-    TestVec,
-    "test_anon_constructor_error_dtype_keyword_missing",
-    test_anon_constructor_error_dtype_keyword_missing,
-    devices=devices,
-)
 add_function_test(
     TestVec,
     "test_anon_constructor_error_length_mismatch",
@@ -1221,20 +1179,14 @@ add_function_test(
 )
 add_function_test(
     TestVec,
-    "test_anon_constructor_error_numeric_arg_missing_1",
-    test_anon_constructor_error_numeric_arg_missing_1,
-    devices=devices,
-)
-add_function_test(
-    TestVec,
-    "test_anon_constructor_error_numeric_arg_missing_2",
-    test_anon_constructor_error_numeric_arg_missing_2,
+    "test_anon_constructor_error_numeric_arg_missing",
+    test_anon_constructor_error_numeric_arg_missing,
     devices=devices,
 )
 add_function_test(
     TestVec,
-    "test_anon_constructor_error_dtype_keyword_extraneous",
-    test_anon_constructor_error_dtype_keyword_extraneous,
+    "test_anon_constructor_error_length_arg_missing",
+    test_anon_constructor_error_length_arg_missing,
     devices=devices,
 )
 add_function_test(
@@ -1258,5 +1210,5 @@ add_function_test(
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2, failfast=True)

warp/tests/test_vec_lite.py

@@ -17,6 +17,7 @@ def test_vector_constructor_value_func():
     b = wp.vector(a, dtype=wp.float16)
     c = wp.vector(a)
     d = wp.vector(a, length=2)
+    e = wp.vector(1.0, length=3)
 
 
 # Test matrix constructors using explicit type (float16)
@@ -67,5 +68,5 @@ add_kernel_test(TestVecLite, test_constructors_default_precision, dim=1, devices
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2, failfast=True)

warp/tests/test_vec_scalar_ops.py

@@ -1902,6 +1902,62 @@ def test_constants(test, device, dtype, register_kernels=False):
     wp.launch(kernel, dim=1, inputs=[], device=device)
 
 
+def test_abs(test, device, dtype, register_kernels=False):
+    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)
+
+    def check_vector_abs():
+        res2 = wp.abs(vec2(wptype(-1), wptype(2)))
+        wp.expect_eq(res2, vec2(wptype(1), wptype(2)))
+
+        res3 = wp.abs(vec3(wptype(1), wptype(-2), wptype(3)))
+        wp.expect_eq(res3, vec3(wptype(1), wptype(2), wptype(3)))
+
+        res4 = wp.abs(vec4(wptype(-1), wptype(2), wptype(3), wptype(-4)))
+        wp.expect_eq(res4, vec4(wptype(1), wptype(2), wptype(3), wptype(4)))
+
+        res5 = wp.abs(vec5(wptype(-1), wptype(2), wptype(-3), wptype(4), wptype(-5)))
+        wp.expect_eq(res5, vec5(wptype(1), wptype(2), wptype(3), wptype(4), wptype(5)))
+
+    kernel = getkernel(check_vector_abs, suffix=dtype.__name__)
+
+    if register_kernels:
+        return
+
+    wp.launch(kernel, dim=1, inputs=[], device=device)
+
+
+def test_sign(test, device, dtype, register_kernels=False):
+    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)
+
+    def check_vector_sign():
+        res2 = wp.sign(vec2(wptype(-1), wptype(2)))
+        wp.expect_eq(res2, vec2(wptype(-1), wptype(1)))
+
+        res3 = wp.sign(vec3(wptype(1), wptype(-2), wptype(3)))
+        wp.expect_eq(res3, vec3(wptype(1), wptype(-1), wptype(1)))
+
+        res4 = wp.sign(vec4(wptype(-1), wptype(2), wptype(3), wptype(-4)))
+        wp.expect_eq(res4, vec4(wptype(-1), wptype(1), wptype(1), wptype(-1)))
+
+        res5 = wp.sign(vec5(wptype(-1), wptype(2), wptype(-3), wptype(4), wptype(-5)))
+        wp.expect_eq(res5, vec5(wptype(-1), wptype(1), wptype(-1), wptype(1), wptype(-1)))
+
+    kernel = getkernel(check_vector_sign, suffix=dtype.__name__)
+
+    if register_kernels:
+        return
+
+    wp.launch(kernel, dim=1, inputs=[], device=device)
+
+
 def test_minmax(test, device, dtype, register_kernels=False):
     rng = np.random.default_rng(123)
 
@@ -2107,10 +2163,18 @@ for dtype in np_scalar_types:
         TestVecScalarOps, f"test_constants_{dtype.__name__}", test_constants, devices=devices, dtype=dtype
     )
 
+    if dtype not in np_unsigned_int_types:
+        add_function_test_register_kernel(
+            TestVecScalarOps, f"test_abs_{dtype.__name__}", test_abs, devices=devices, dtype=dtype
+        )
+        add_function_test_register_kernel(
+            TestVecScalarOps, f"test_sign_{dtype.__name__}", test_sign, devices=devices, dtype=dtype
+        )
+
     # the kernels in this test compile incredibly slowly...
     # add_function_test_register_kernel(TestVecScalarOps, f"test_minmax_{dtype.__name__}", test_minmax, devices=devices, dtype=dtype)
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2, failfast=True)

warp/tests/test_verify_fp.py

@@ -88,5 +88,5 @@ add_function_test(TestVerifyFP, "test_nan", test_nan, devices=devices, check_out
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_volume.py

@@ -843,8 +843,31 @@ def test_volume_from_numpy(test, device):
     test.assertIsNone(sphere_vdb_array.deleter)
 
 
+def test_volume_aniso_transform(test, device):
+    # XY-rotation + z scale
+    transform = [
+        [0, -1, 0],
+        [1, 0, 0],
+        [0, 0, 2],
+    ]
+
+    points = wp.array([[-1, 1, 4]], dtype=float, device=device)
+    volume = wp.Volume.allocate_by_voxels(voxel_points=points, transform=transform, device=device)
+
+    # Check that world points are correctly converted to local space
+    voxels = volume.get_voxels().numpy()
+    assert_np_equal(voxels, [[1, 1, 2]])
+
+    # Check that we retrieve the correct transform from the grid metadata
+    assert_np_equal(volume.get_voxel_size(), [-1, 1, 2])
+    assert_np_equal(transform, np.array(volume.get_grid_info().transform_matrix).reshape(3, 3))
+
+
 class TestVolume(unittest.TestCase):
-    pass
+    def test_volume_new_del(self):
+        # test the scenario in which a volume is created but not initialized before gc
+        instance = wp.Volume.__new__(wp.Volume)
+        instance.__del__()
 
 
 add_function_test(
@@ -871,6 +894,9 @@ add_function_test(TestVolume, "test_volume_introspection", test_volume_introspec
 add_function_test(
     TestVolume, "test_volume_from_numpy", test_volume_from_numpy, devices=get_selected_cuda_test_devices()
 )
+add_function_test(
+    TestVolume, "test_volume_aniso_transform", test_volume_aniso_transform, devices=get_selected_cuda_test_devices()
+)
 add_function_test(TestVolume, "test_volume_multiple_grids", test_volume_multiple_grids, devices=devices)
 add_function_test(TestVolume, "test_volume_feature_array", test_volume_feature_array, devices=devices)
 add_function_test(TestVolume, "test_volume_sample_index", test_volume_sample_index, devices=devices)
@@ -957,5 +983,5 @@ for device in devices:
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_volume_write.py

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

warp/tests/unittest_serial.py

@@ -15,7 +15,7 @@ def run_suite() -> bool:
     """Run a test suite"""
 
     # force rebuild of all kernels
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     print("Cleared Warp kernel cache")
 
     runner = unittest.TextTestRunner(verbosity=2, failfast=True)

warp/tests/unittest_suites.py

@@ -120,6 +120,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
     from warp.tests.test_generics import TestGenerics
     from warp.tests.test_grad import TestGrad
     from warp.tests.test_grad_customs import TestGradCustoms
+    from warp.tests.test_grad_debug import TestGradDebug
     from warp.tests.test_hash_grid import TestHashGrid
     from warp.tests.test_import import TestImport
     from warp.tests.test_indexedarray import TestIndexedArray
@@ -150,6 +151,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
     from warp.tests.test_noise import TestNoise
     from warp.tests.test_operators import TestOperators
     from warp.tests.test_options import TestOptions
+    from warp.tests.test_overwrite import TestOverwrite
     from warp.tests.test_peer import TestPeer
     from warp.tests.test_pinned import TestPinned
     from warp.tests.test_print import TestPrint
@@ -164,6 +166,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
     from warp.tests.test_snippet import TestSnippets
     from warp.tests.test_sparse import TestSparse
     from warp.tests.test_spatial import TestSpatial
+    from warp.tests.test_special_values import TestSpecialValues
     from warp.tests.test_streams import TestStreams
     from warp.tests.test_struct import TestStruct
     from warp.tests.test_tape import TestTape
@@ -211,6 +214,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestGenerics,
         TestGrad,
         TestGradCustoms,
+        TestGradDebug,
         TestHashGrid,
         TestImport,
         TestIndexedArray,
@@ -241,6 +245,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestNoise,
         TestOperators,
         TestOptions,
+        TestOverwrite,
         TestPeer,
         TestPinned,
         TestPrint,
@@ -252,9 +257,10 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestSimGradients,
         TestSimKinematics,
         TestSmoothstep,
-        TestSparse,
         TestSnippets,
+        TestSparse,
         TestSpatial,
+        TestSpecialValues,
         TestStreams,
         TestStruct,
         TestTape,
@@ -291,6 +297,7 @@ def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
     from warp.tests.test_func import TestFunc
     from warp.tests.test_generics import TestGenerics
     from warp.tests.test_grad_customs import TestGradCustoms
+    from warp.tests.test_grad_debug import TestGradDebug
     from warp.tests.test_hash_grid import TestHashGrid
     from warp.tests.test_indexedarray import TestIndexedArray
     from warp.tests.test_launch import TestLaunch
@@ -334,6 +341,7 @@ def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
         TestFunc,
         TestGenerics,
         TestGradCustoms,
+        TestGradDebug,
         TestHashGrid,
         TestIndexedArray,
         TestLaunch,

warp/tests/walkthrough_debug.py

@@ -81,5 +81,5 @@ def example_breakpoint(n: int):
 
 
 if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     wp.launch(example_breakpoint, dim=1, inputs=[10], device="cpu")

warp/thirdparty/unittest_parallel.py

@@ -173,7 +173,7 @@ def main(argv=None):
     import warp as wp  # NVIDIA Modification
 
     # Clear the Warp cache (NVIDIA Modification)
-    wp.build.clear_kernel_cache()
+    wp.clear_kernel_cache()
     print("Cleared Warp kernel cache")
 
     # Create the temporary directory (for coverage files)
@@ -244,7 +244,7 @@ def main(argv=None):
                     initargs=(manager.Lock(), shared_index, args, temp_dir),
                 ) as executor:
                     test_manager = ParallelTestManager(manager, args, temp_dir)
-                    results = list(executor.map(test_manager.run_tests, test_suites, timeout=7200))
+                    results = list(executor.map(test_manager.run_tests, test_suites, timeout=2400))
         else:
             # This entire path is an NVIDIA Modification
 
@@ -554,9 +554,6 @@ def initialize_test_process(lock, shared_index, args, temp_dir):
             wp.config.kernel_cache_dir = cache_root_dir
 
             wp.build.clear_kernel_cache()
-        else:
-            # Initialize Warp is if hasn't been initialized already
-            wp.init()
 
 
 if __name__ == "__main__":  # pragma: no cover

warp/torch.py

@@ -10,12 +10,41 @@ import ctypes
 import numpy
 
 import warp
+import warp.context
 
 
 # return the warp device corresponding to a torch device
 def device_from_torch(torch_device) -> warp.context.Device:
-    """Return the Warp device corresponding to a Torch device."""
-    return warp.get_device(str(torch_device))
+    """Return the Warp device corresponding to a Torch device.
+
+    Args:
+        torch_device (`torch.device` or `str`): Torch device identifier
+
+    Raises:
+        RuntimeError: Torch device does not have a corresponding Warp device
+    """
+    if type(torch_device) is str:
+        warp_device = warp.context.runtime.device_map.get(torch_device)
+        if warp_device is not None:
+            return warp_device
+        elif torch_device == "cuda":
+            return warp.context.runtime.get_current_cuda_device()
+        else:
+            raise RuntimeError(f"Unsupported Torch device {torch_device}")
+    else:
+        try:
+            if torch_device.type == "cuda":
+                return warp.context.runtime.cuda_devices[torch_device.index]
+            elif torch_device.type == "cpu":
+                return warp.context.runtime.cpu_device
+            else:
+                raise RuntimeError(f"Unsupported Torch device type {torch_device.type}")
+        except Exception as e:
+            import torch
+
+            if not isinstance(torch_device, torch.device):
+                raise ValueError("Argument must be a torch.device object or a string") from e
+            raise
 
 
 def device_to_torch(warp_device: warp.context.Devicelike) -> str:
@@ -154,16 +183,17 @@ dtype_is_compatible.compatible_sets = None
 
 
 # wrap a torch tensor to a wp array, data is not copied
-def from_torch(t, dtype=None, requires_grad=None, grad=None):
+def from_torch(t, dtype=None, requires_grad=None, grad=None, return_ctype=False):
     """Convert a Torch tensor to a Warp array without copying the data.
 
     Args:
         t (torch.Tensor): The torch tensor to wrap.
         dtype (warp.dtype, optional): The target data type of the resulting Warp array. Defaults to the tensor value type mapped to a Warp array value type.
         requires_grad (bool, optional): Whether the resulting array should wrap the tensor's gradient, if it exists (the grad tensor will be allocated otherwise). Defaults to the tensor's `requires_grad` value.
+        return_ctype (bool, optional): Whether to return a low-level array descriptor instead of a ``wp.array`` object (faster).  The descriptor can be passed to Warp kernels.
 
     Returns:
-        warp.array: The wrapped array.
+        warp.array: The wrapped array or array descriptor.
     """
     if dtype is None:
         dtype = dtype_from_torch(t.dtype)
@@ -175,7 +205,6 @@ def from_torch(t, dtype=None, requires_grad=None, grad=None):
 
     shape = tuple(t.shape)
     strides = tuple(s * ctype_size for s in t.stride())
-    device = device_from_torch(t.device)
 
     # if target is a vector or matrix type
     # then check if trailing dimensions match
@@ -183,57 +212,90 @@ def from_torch(t, dtype=None, requires_grad=None, grad=None):
     if hasattr(dtype, "_shape_"):
         dtype_shape = dtype._shape_
         dtype_dims = len(dtype._shape_)
+        # ensure inner shape matches
         if dtype_dims > len(shape) or dtype_shape != shape[-dtype_dims:]:
             raise RuntimeError(
                 f"Could not convert Torch tensor with shape {shape} to Warp array with dtype={dtype}, ensure that source inner shape is {dtype_shape}"
             )
-
-        # ensure the inner strides are contiguous
-        stride = ctype_size
-        for i in range(dtype_dims):
-            if strides[-i - 1] != stride:
-                raise RuntimeError(
-                    f"Could not convert Torch tensor with shape {shape} to Warp array with dtype={dtype}, because the source inner strides are not contiguous"
-                )
-            stride *= dtype_shape[-i - 1]
-
+        # ensure inner strides are contiguous
+        if strides[-1] != ctype_size or (dtype_dims > 1 and strides[-2] != ctype_size * dtype_shape[-1]):
+            raise RuntimeError(
+                f"Could not convert Torch tensor with shape {shape} to Warp array with dtype={dtype}, because the source inner strides are not contiguous"
+            )
+        # trim shape and strides
         shape = tuple(shape[:-dtype_dims]) or (1,)
         strides = tuple(strides[:-dtype_dims]) or (ctype_size,)
 
+    # gradient
+    # - if return_ctype is False, we set `grad` to a wp.array or None
+    # - if return_ctype is True, we set `grad_ptr` and set `grad` as the owner (wp.array or torch.Tensor)
     requires_grad = t.requires_grad if requires_grad is None else requires_grad
+    grad_ptr = 0
     if grad is not None:
-        if not isinstance(grad, warp.array):
-            import torch
-
-            if isinstance(grad, torch.Tensor):
-                grad = from_torch(grad, dtype=dtype)
+        if isinstance(grad, warp.array):
+            if return_ctype:
+                if grad.strides != strides:
+                    raise RuntimeError(
+                        f"Gradient strides must match array strides, expected {strides} but got {grad.strides}"
+                    )
+                grad_ptr = grad.ptr
+        else:
+            # assume grad is a torch.Tensor
+            if return_ctype:
+                if t.stride() != grad.stride():
+                    raise RuntimeError(
+                        f"Gradient strides must match array strides, expected {t.stride()} but got {grad.stride()}"
+                    )
+                grad_ptr = grad.data_ptr()
             else:
-                raise ValueError(f"Invalid gradient type: {type(grad)}")
+                grad = from_torch(grad, dtype=dtype, requires_grad=False)
     elif requires_grad:
         # wrap the tensor gradient, allocate if necessary
-        if t.grad is None:
+        if t.grad is not None:
+            if return_ctype:
+                grad = t.grad
+                if t.stride() != grad.stride():
+                    raise RuntimeError(
+                        f"Gradient strides must match array strides, expected {t.stride()} but got {grad.stride()}"
+                    )
+                grad_ptr = grad.data_ptr()
+            else:
+                grad = from_torch(t.grad, dtype=dtype, requires_grad=False)
+        else:
             # allocate a zero-filled gradient if it doesn't exist
             # Note: we use Warp to allocate the shared gradient with compatible strides
-            grad = warp.zeros(dtype=dtype, shape=shape, strides=strides, device=device)
+            grad = warp.zeros(dtype=dtype, shape=shape, strides=strides, device=device_from_torch(t.device))
             t.grad = to_torch(grad, requires_grad=False)
-        else:
-            # TODO: this will fail if the strides are incompatible
-            grad = from_torch(t.grad, dtype=dtype)
-
-    a = warp.array(
-        ptr=t.data_ptr(),
-        dtype=dtype,
-        shape=shape,
-        strides=strides,
-        device=device,
-        copy=False,
-        grad=grad,
-        requires_grad=requires_grad,
-    )
-
-    # save a reference to the source tensor, otherwise it will be deallocated
-    a._tensor = t
-    return a
+            grad_ptr = grad.ptr
+
+    if return_ctype:
+        ptr = t.data_ptr()
+
+        # create array descriptor
+        array_ctype = warp.types.array_t(ptr, grad_ptr, len(shape), shape, strides)
+
+        # keep data and gradient alive
+        array_ctype._ref = t
+        array_ctype._gradref = grad
+
+        return array_ctype
+
+    else:
+        a = warp.array(
+            ptr=t.data_ptr(),
+            dtype=dtype,
+            shape=shape,
+            strides=strides,
+            device=device_from_torch(t.device),
+            copy=False,
+            grad=grad,
+            requires_grad=requires_grad,
+        )
+
+        # save a reference to the source tensor, otherwise it may get deallocated
+        a._tensor = t
+
+        return a
 
 
 def to_torch(a, requires_grad=None):