warp-lang 1.1.0__py3-none-manylinux2014_x86_64.whl → 1.2.0__py3-none-manylinux2014_x86_64.whl

warp/tests/test_examples.py

@@ -27,7 +27,7 @@ import os
 import subprocess
 import sys
 import unittest
-from typing import Any, Dict
+from typing import Any, Dict, Optional, Type
 
 import warp as wp
 import warp.tests.unittest_utils
@@ -38,10 +38,9 @@ from warp.tests.unittest_utils import (
     sanitize_identifier,
 )
 
-wp.init()
-
 
 def _build_command_line_options(test_options: Dict[str, Any]) -> list:
+    """Helper function to build command-line options from the test options dictionary."""
     additional_options = []
 
     for key, value in test_options.items():
@@ -54,8 +53,25 @@ def _build_command_line_options(test_options: Dict[str, Any]) -> list:
     return additional_options
 
 
-# registers an example to run as a TestCase
-def add_example_test(cls, name, devices=None, test_options=None, test_options_cpu=None, test_options_cuda=None):
+def _merge_options(base_options: Dict[str, Any], device_options: Dict[str, Any]) -> Dict[str, Any]:
+    """Helper function to merge base test options with device-specific test options."""
+    merged_options = base_options.copy()
+
+    #  Update options with device-specific dictionary, overwriting existing keys with the more-specific values
+    merged_options.update(device_options)
+    return merged_options
+
+
+def add_example_test(
+    cls: Type,
+    name: str,
+    devices: Optional[list] = None,
+    test_options: Optional[Dict[str, Any]] = None,
+    test_options_cpu: Optional[Dict[str, Any]] = None,
+    test_options_cuda: Optional[Dict[str, Any]] = None,
+):
+    """Registers a Warp example to run on ``devices`` as a TestCase."""
+
     if test_options is None:
         test_options = {}
     if test_options_cpu is None:
@@ -64,7 +80,11 @@ def add_example_test(cls, name, devices=None, test_options=None, test_options_cp
         test_options_cuda = {}
 
     def run(test, device):
-        options = test_options.copy()
+        if wp.get_device(device).is_cuda:
+            options = _merge_options(test_options, test_options_cuda)
+        else:
+            options = _merge_options(test_options, test_options_cpu)
+
         # Mark the test as skipped if Torch is not installed but required
         torch_required = options.pop("torch_required", False)
         if torch_required:
@@ -99,11 +119,7 @@ def add_example_test(cls, name, devices=None, test_options=None, test_options_cp
             ) as coverage_file:
                 pass
 
-            command = [
-                "coverage",
-                "run",
-                f"--data-file={coverage_file.name}",
-            ]
+            command = ["coverage", "run", f"--data-file={coverage_file.name}"]
 
             if warp.tests.unittest_utils.coverage_branch:
                 command.append("--branch")
@@ -132,13 +148,6 @@ def add_example_test(cls, name, devices=None, test_options=None, test_options_cp
 
         command.extend(_build_command_line_options(options))
 
-        if wp.get_device(device).is_cuda:
-            # Overrides when testing on a CUDA device
-            command.extend(_build_command_line_options(test_options_cuda))
-        else:
-            # Overrides when testing on a CPU device
-            command.extend(_build_command_line_options(test_options_cpu))
-
         # with wp.ScopedTimer(f"{name}_{sanitize_identifier(device)}"):
         # Run the script as a subprocess
         result = subprocess.run(command, capture_output=True, text=True, env=env_vars)
@@ -176,12 +185,7 @@ class TestCoreExamples(unittest.TestCase):
 # Exclude unless we can run headless somehow
 # add_example_test(TestCoreExamples, name="example_render_opengl")
 
-add_example_test(
-    TestCoreExamples,
-    name="core.example_dem",
-    devices=test_devices,
-    test_options_cpu={"num_frames": 2},
-)
+add_example_test(TestCoreExamples, name="core.example_dem", devices=test_devices, test_options_cpu={"num_frames": 2})
 add_example_test(
     TestCoreExamples,
     name="core.example_fluid",
@@ -261,12 +265,16 @@ add_example_test(
     devices=test_devices,
     test_options={"headless": True, "train_iters": 50},
 )
+# NOTE: This example uses CUTLASS and will run orders of magnitude slower when Warp is built in debug mode
 add_example_test(
     TestOptimExamples,
     name="optim.example_walker",
     devices=test_devices,
     test_options={"usd_required": True},
-    test_options_cuda={"train_iters": 3, "num_frames": 60},
+    test_options_cuda={
+        "train_iters": 1 if warp.context.runtime.core.is_debug_enabled() else 3,
+        "num_frames": 1 if warp.context.runtime.core.is_debug_enabled() else 60,
+    },
     test_options_cpu={"train_iters": 1, "num_frames": 30},
 )
 
@@ -301,11 +309,8 @@ add_example_test(
     test_options_cpu={"num_frames": 3},
 )
 add_example_test(
-    TestSimExamples,
-    name="sim.example_rigid_soft_contact",
-    devices=cuda_test_devices,
-    test_options_cpu={"num_frames": 10},
-)  # TODO: Segfaults on CPU in some cases, add back in CPU testing when fixed
+    TestSimExamples, name="sim.example_rigid_soft_contact", devices=test_devices, test_options_cpu={"num_frames": 10}
+)
 add_example_test(TestSimExamples, name="sim.example_rigid_force", devices=test_devices)
 add_example_test(TestSimExamples, name="sim.example_rigid_gyroscopic", devices=test_devices)
 add_example_test(
@@ -328,13 +333,14 @@ add_example_test(
     test_options={"headless": True},
 )
 
-# The following examples do not need CUDA
 add_example_test(
     TestFemExamples,
     name="fem.example_apic_fluid",
-    devices=test_devices,
-    test_options={"num_frames": 1, "res": "16,16,16"},
+    devices=get_selected_cuda_test_devices(),
+    test_options={"num_frames": 5, "voxel_size": 2.0},
 )
+
+# The following examples do not need CUDA
 add_example_test(
     TestFemDiffusionExamples,
     name="fem.example_diffusion",
@@ -342,10 +348,7 @@ add_example_test(
     test_options={"resolution": 10, "mesh": "tri", "headless": True},
 )
 add_example_test(
-    TestFemDiffusionExamples,
-    name="fem.example_diffusion_3d",
-    devices=test_devices,
-    test_options={"headless": True},
+    TestFemDiffusionExamples, name="fem.example_diffusion_3d", devices=test_devices, test_options={"headless": True}
 )
 add_example_test(
     TestFemExamples,
@@ -361,9 +364,9 @@ add_example_test(
 )
 add_example_test(
     TestFemExamples,
-    name="fem.example_convection_diffusion_dg0",
+    name="fem.example_burgers",
     devices=test_devices,
-    test_options={"resolution": 20, "num_frames": 25, "mesh": "quad", "headless": True},
+    test_options={"resolution": 20, "num_frames": 25, "degree": 1, "headless": True},
 )
 add_example_test(
     TestFemExamples,

warp/tests/test_fabricarray.py

@@ -14,9 +14,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
-
 # types to test fabric arrays
 _fabric_types = [
     *wp.types.scalar_types,

warp/tests/test_fast_math.py

@@ -10,8 +10,6 @@ import unittest
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def test_pow(e: float, result: float):

warp/tests/test_fem.py

@@ -21,8 +21,6 @@ from warp.fem.types import make_free_sample
 from warp.fem.utils import grid_to_hexes, grid_to_quads, grid_to_tets, grid_to_tris
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @integrand
 def linear_form(s: Sample, u: Field):
@@ -435,13 +433,13 @@ def _launch_test_geometry_kernel(geo: fem.Geometry, device):
         wp.expect_near(coords, inner_side_coords, 0.0001)
         wp.expect_near(coords, outer_side_coords, 0.0001)
 
-        vol = geo.side_measure(side_arg, s)
-        wp.atomic_add(side_measures, side_index, vol * qp_weights[q])
+        area = geo.side_measure(side_arg, s)
+        wp.atomic_add(side_measures, side_index, area * qp_weights[q])
 
         # test consistency of side normal, measure, and deformation gradient
         F = geo.side_deformation_gradient(side_arg, s)
         F_det = DeformedGeometry._side_measure(F)
-        wp.expect_near(F_det * REF_MEASURE, vol)
+        wp.expect_near(F_det * REF_MEASURE, area)
 
         nor = geo.side_normal(side_arg, s)
         F_cross = DeformedGeometry._side_normal(F)
@@ -589,6 +587,28 @@ def test_hex_mesh(test, device):
     assert_np_equal(cell_measures.numpy(), np.full(cell_measures.shape, 1.0 / (N**3)), tol=1.0e-4)
 
 
+def test_nanogrid(test, device):
+    N = 8
+
+    points = wp.array([[0.5, 0.5, 0.5]], dtype=float, device=device)
+    volume = wp.Volume.allocate_by_tiles(
+        tile_points=points, voxel_size=1.0 / N, translation=(0.0, 0.0, 0.0), bg_value=None, device=device
+    )
+
+    geo = fem.Nanogrid(volume)
+
+    test.assertEqual(geo.cell_count(), (N) ** 3)
+    test.assertEqual(geo.vertex_count(), (N + 1) ** 3)
+    test.assertEqual(geo.side_count(), 3 * (N + 1) * N**2)
+    test.assertEqual(geo.boundary_side_count(), 6 * N * N)
+    test.assertEqual(geo.edge_count(), 3 * N * (N + 1) ** 2)
+
+    side_measures, cell_measures = _launch_test_geometry_kernel(geo, device)
+
+    assert_np_equal(side_measures.numpy(), np.full(side_measures.shape, 1.0 / (N**2)), tol=1.0e-4)
+    assert_np_equal(cell_measures.numpy(), np.full(cell_measures.shape, 1.0 / (N**3)), tol=1.0e-4)
+
+
 @integrand
 def _rigid_deformation_field(s: Sample, domain: Domain, translation: wp.vec3, rotation: wp.vec3, scale: float):
     q = wp.quat_from_axis_angle(wp.normalize(rotation), wp.length(rotation))
@@ -601,63 +621,62 @@ def test_deformed_geometry(test, device):
     with wp.ScopedDevice(device):
         positions, tet_vidx = _gen_tetmesh(N)
 
-    geo = fem.Tetmesh(tet_vertex_indices=tet_vidx, positions=positions)
+        geo = fem.Tetmesh(tet_vertex_indices=tet_vidx, positions=positions)
 
-    translation = [1.0, 2.0, 3.0]
-    rotation = [0.0, math.pi / 4.0, 0.0]
-    scale = 2.0
+        translation = [1.0, 2.0, 3.0]
+        rotation = [0.0, math.pi / 4.0, 0.0]
+        scale = 2.0
 
-    vector_space = fem.make_polynomial_space(geo, dtype=wp.vec3, degree=2)
-    pos_field = vector_space.make_field()
-    fem.interpolate(
-        _rigid_deformation_field,
-        dest=pos_field,
-        values={"translation": translation, "rotation": rotation, "scale": scale},
-    )
-
-    deformed_geo = pos_field.make_deformed_geometry()
-
-    # rigidly-deformed geometry
-
-    test.assertEqual(geo.cell_count(), 5 * (N) ** 3)
-    test.assertEqual(geo.vertex_count(), (N + 1) ** 3)
-    test.assertEqual(geo.side_count(), 6 * (N + 1) * N**2 + (N**3) * 4)
-    test.assertEqual(geo.boundary_side_count(), 12 * N * N)
+        vector_space = fem.make_polynomial_space(geo, dtype=wp.vec3, degree=2)
+        pos_field = vector_space.make_field()
+        fem.interpolate(
+            _rigid_deformation_field,
+            dest=pos_field,
+            values={"translation": translation, "rotation": rotation, "scale": scale},
+        )
 
-    side_measures, cell_measures = _launch_test_geometry_kernel(deformed_geo, device)
+        deformed_geo = pos_field.make_deformed_geometry()
 
-    test.assertAlmostEqual(
-        np.sum(cell_measures.numpy()), scale**3, places=4, msg=f"cell_measures = {cell_measures.numpy()}"
-    )
-    test.assertAlmostEqual(
-        np.sum(side_measures.numpy()), scale**2 * (0.5 * 6 * (N + 1) + N * 2 * math.sqrt(3.0)), places=4
-    )
+        # rigidly-deformed geometry
 
-    @wp.kernel
-    def _test_deformed_geometry_normal(
-        geo_index_arg: geo.SideIndexArg, geo_arg: geo.SideArg, def_arg: deformed_geo.SideArg, rotation: wp.vec3
-    ):
-        i = wp.tid()
-        side_index = deformed_geo.boundary_side_index(geo_index_arg, i)
+        test.assertEqual(geo.cell_count(), 5 * (N) ** 3)
+        test.assertEqual(geo.vertex_count(), (N + 1) ** 3)
+        test.assertEqual(geo.side_count(), 6 * (N + 1) * N**2 + (N**3) * 4)
+        test.assertEqual(geo.boundary_side_count(), 12 * N * N)
 
-        s = make_free_sample(side_index, Coords(0.5, 0.5, 0.0))
-        geo_n = geo.side_normal(geo_arg, s)
-        def_n = deformed_geo.side_normal(def_arg, s)
+        side_measures, cell_measures = _launch_test_geometry_kernel(deformed_geo, wp.get_device())
 
-        q = wp.quat_from_axis_angle(wp.normalize(rotation), wp.length(rotation))
-        wp.expect_near(wp.quat_rotate(q, geo_n), def_n, 0.001)
+        test.assertAlmostEqual(
+            np.sum(cell_measures.numpy()), scale**3, places=4, msg=f"cell_measures = {cell_measures.numpy()}"
+        )
+        test.assertAlmostEqual(
+            np.sum(side_measures.numpy()), scale**2 * (0.5 * 6 * (N + 1) + N * 2 * math.sqrt(3.0)), places=4
+        )
 
-    wp.launch(
-        _test_deformed_geometry_normal,
-        dim=geo.boundary_side_count(),
-        device=device,
-        inputs=[
-            geo.side_index_arg_value(device),
-            geo.side_arg_value(device),
-            deformed_geo.side_arg_value(device),
-            rotation,
-        ],
-    )
+        @wp.kernel
+        def _test_deformed_geometry_normal(
+            geo_index_arg: geo.SideIndexArg, geo_arg: geo.SideArg, def_arg: deformed_geo.SideArg, rotation: wp.vec3
+        ):
+            i = wp.tid()
+            side_index = deformed_geo.boundary_side_index(geo_index_arg, i)
+
+            s = make_free_sample(side_index, Coords(0.5, 0.5, 0.0))
+            geo_n = geo.side_normal(geo_arg, s)
+            def_n = deformed_geo.side_normal(def_arg, s)
+
+            q = wp.quat_from_axis_angle(wp.normalize(rotation), wp.length(rotation))
+            wp.expect_near(wp.quat_rotate(q, geo_n), def_n, 0.001)
+
+        wp.launch(
+            _test_deformed_geometry_normal,
+            dim=geo.boundary_side_count(),
+            inputs=[
+                geo.side_index_arg_value(wp.get_device()),
+                geo.side_arg_value(wp.get_device()),
+                deformed_geo.side_arg_value(wp.get_device()),
+                rotation,
+            ],
+        )
 
     wp.synchronize()
 
@@ -1237,6 +1256,7 @@ def test_particle_quadratures(test, device):
 
 
 devices = get_test_devices()
+cuda_devices = get_selected_cuda_test_devices()
 
 
 class TestFem(unittest.TestCase):
@@ -1256,6 +1276,7 @@ add_function_test(TestFem, "test_quad_mesh", test_quad_mesh, devices=devices)
 add_function_test(TestFem, "test_grid_3d", test_grid_3d, devices=devices)
 add_function_test(TestFem, "test_tet_mesh", test_tet_mesh, devices=devices)
 add_function_test(TestFem, "test_hex_mesh", test_hex_mesh, devices=devices)
+add_function_test(TestFem, "test_nanogrid", test_nanogrid, devices=cuda_devices)
 add_function_test(TestFem, "test_deformed_geometry", test_deformed_geometry, devices=devices)
 add_function_test(TestFem, "test_dof_mapper", test_dof_mapper)
 add_function_test(TestFem, "test_point_basis", test_point_basis)

warp/tests/test_fp16.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def load_store_half(f32: wp.array(dtype=wp.float32), f16: wp.array(dtype=wp.float16)):

warp/tests/test_func.py

@@ -13,8 +13,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.func
 def sqr(x: float):

warp/tests/test_generics.py

@@ -13,8 +13,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.func
 def generic_adder(a: Any, b: Any):
@@ -240,6 +238,32 @@ def test_generic_fill_overloads(test, device):
         assert_np_equal(a3b.numpy(), np.full((n, 3), True, dtype=np.bool_))
 
 
+# generic kernel used to test generic types mixed with specialized types
+@wp.func
+def generic_conditional_setter_func(a: wp.array(dtype=Any), i: int, value: Any, relative: bool):
+    if relative:
+        a[i] += value
+    else:
+        a[i] = value
+
+
+@wp.kernel
+def generic_conditional_setter(a: wp.array(dtype=Any), i: int, value: Any, relative: bool):
+    generic_conditional_setter_func(a, i, value, relative)
+
+
+def test_generic_conditional_setter(test, device):
+    with wp.ScopedDevice(device):
+        n = 10
+        ai = wp.zeros(n, dtype=int)
+
+        wp.launch(generic_conditional_setter, dim=1, inputs=[ai, 1, 42, False])
+        wp.launch(generic_conditional_setter, dim=1, inputs=[ai, 1, 5, True])
+        ai_true = np.zeros(n, dtype=np.int32)
+        ai_true[1] = 47
+        assert_np_equal(ai.numpy(), ai_true)
+
+
 # custom vector/matrix types
 my_vec5 = wp.vec(length=5, dtype=wp.float32)
 my_mat55 = wp.mat(shape=(5, 5), dtype=wp.float32)
@@ -511,6 +535,7 @@ add_function_test(TestGenerics, "test_generic_array_kernel", test_generic_array_
 add_function_test(TestGenerics, "test_generic_accumulator_kernel", test_generic_accumulator_kernel, devices=devices)
 add_function_test(TestGenerics, "test_generic_fill", test_generic_fill, devices=devices)
 add_function_test(TestGenerics, "test_generic_fill_overloads", test_generic_fill_overloads, devices=devices)
+add_function_test(TestGenerics, "test_generic_conditional_setter", test_generic_conditional_setter, devices=devices)
 add_function_test(TestGenerics, "test_generic_transform_kernel", test_generic_transform_kernel, devices=devices)
 add_function_test(
     TestGenerics, "test_generic_transform_array_kernel", test_generic_transform_array_kernel, devices=devices

warp/tests/test_grad.py

@@ -13,8 +13,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def scalar_grad(x: wp.array(dtype=float), y: wp.array(dtype=float)):
@@ -665,16 +663,40 @@ def square_kernel(x: wp.array(dtype=float), y: wp.array(dtype=float)):
     y[tid] = x[tid] ** 2.0
 
 
+@wp.kernel
+def square_slice_2d_kernel(x: wp.array2d(dtype=float), y: wp.array2d(dtype=float), row_idx: int):
+    tid = wp.tid()
+    x_slice = x[row_idx]
+    y_slice = y[row_idx]
+    y_slice[tid] = x_slice[tid] ** 2.0
+
+
+@wp.kernel
+def square_slice_3d_1d_kernel(x: wp.array3d(dtype=float), y: wp.array3d(dtype=float), slice_idx: int):
+    i, j = wp.tid()
+    x_slice = x[slice_idx]
+    y_slice = y[slice_idx]
+    y_slice[i, j] = x_slice[i, j] ** 2.0
+
+
+@wp.kernel
+def square_slice_3d_2d_kernel(x: wp.array3d(dtype=float), y: wp.array3d(dtype=float), slice_i: int, slice_j: int):
+    tid = wp.tid()
+    x_slice = x[slice_i, slice_j]
+    y_slice = y[slice_i, slice_j]
+    y_slice[tid] = x_slice[tid] ** 2.0
+
+
 def test_gradient_internal(test, device):
     with wp.ScopedDevice(device):
         a = wp.array([1.0, 2.0, 3.0], dtype=float, requires_grad=True)
         b = wp.array([0.0, 0.0, 0.0], dtype=float, requires_grad=True)
 
-        wp.launch(square_kernel, a.size, inputs=[a, b])
+        wp.launch(square_kernel, dim=a.size, inputs=[a, b])
 
         # use internal gradients (.grad), adj_inputs are None
         b.grad = wp.array([1.0, 1.0, 1.0], dtype=float)
-        wp.launch(square_kernel, a.shape[0], inputs=[a, b], adjoint=True, adj_inputs=[None, None])
+        wp.launch(square_kernel, dim=a.size, inputs=[a, b], adjoint=True, adj_inputs=[None, None])
 
         assert_np_equal(a.grad.numpy(), np.array([2.0, 4.0, 6.0]))
 
@@ -684,12 +706,12 @@ def test_gradient_external(test, device):
         a = wp.array([1.0, 2.0, 3.0], dtype=float, requires_grad=False)
         b = wp.array([0.0, 0.0, 0.0], dtype=float, requires_grad=False)
 
-        wp.launch(square_kernel, a.size, inputs=[a, b])
+        wp.launch(square_kernel, dim=a.size, inputs=[a, b])
 
         # use external gradients passed in adj_inputs
         a_grad = wp.array([0.0, 0.0, 0.0], dtype=float)
         b_grad = wp.array([1.0, 1.0, 1.0], dtype=float)
-        wp.launch(square_kernel, a.shape[0], inputs=[a, b], adjoint=True, adj_inputs=[a_grad, b_grad])
+        wp.launch(square_kernel, dim=a.size, inputs=[a, b], adjoint=True, adj_inputs=[a_grad, b_grad])
 
         assert_np_equal(a_grad.numpy(), np.array([2.0, 4.0, 6.0]))
 
@@ -699,18 +721,132 @@ def test_gradient_precedence(test, device):
         a = wp.array([1.0, 2.0, 3.0], dtype=float, requires_grad=True)
         b = wp.array([0.0, 0.0, 0.0], dtype=float, requires_grad=True)
 
-        wp.launch(square_kernel, a.size, inputs=[a, b])
+        wp.launch(square_kernel, dim=a.size, inputs=[a, b])
 
         # if both internal and external gradients are present, the external one takes precedence,
         # because it's explicitly passed by the user in adj_inputs
         a_grad = wp.array([0.0, 0.0, 0.0], dtype=float)
         b_grad = wp.array([1.0, 1.0, 1.0], dtype=float)
-        wp.launch(square_kernel, a.shape[0], inputs=[a, b], adjoint=True, adj_inputs=[a_grad, b_grad])
+        wp.launch(square_kernel, dim=a.size, inputs=[a, b], adjoint=True, adj_inputs=[a_grad, b_grad])
 
         assert_np_equal(a_grad.numpy(), np.array([2.0, 4.0, 6.0]))  # used
         assert_np_equal(a.grad.numpy(), np.array([0.0, 0.0, 0.0]))  # unused
 
 
+def test_gradient_slice_2d(test, device):
+    with wp.ScopedDevice(device):
+        a = wp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=float, requires_grad=True)
+        b = wp.zeros_like(a, requires_grad=False)
+        b.grad = wp.ones_like(a, requires_grad=False)
+
+        wp.launch(square_slice_2d_kernel, dim=a.shape[1], inputs=[a, b, 1])
+
+        # use internal gradients (.grad), adj_inputs are None
+        wp.launch(square_slice_2d_kernel, dim=a.shape[1], inputs=[a, b, 1], adjoint=True, adj_inputs=[None, None, 1])
+
+        assert_np_equal(a.grad.numpy(), np.array([[0.0, 0.0], [6.0, 8.0], [0.0, 0.0]]))
+
+
+def test_gradient_slice_3d_1d(test, device):
+    with wp.ScopedDevice(device):
+        data = [
+            [
+                [1, 2, 3],
+                [4, 5, 6],
+                [7, 8, 9],
+            ],
+            [
+                [11, 12, 13],
+                [14, 15, 16],
+                [17, 18, 19],
+            ],
+            [
+                [21, 22, 23],
+                [24, 25, 26],
+                [27, 28, 29],
+            ],
+        ]
+        a = wp.array(data, dtype=float, requires_grad=True)
+        b = wp.zeros_like(a, requires_grad=False)
+        b.grad = wp.ones_like(a, requires_grad=False)
+
+        wp.launch(square_slice_3d_1d_kernel, dim=a.shape[1:], inputs=[a, b, 1])
+
+        # use internal gradients (.grad), adj_inputs are None
+        wp.launch(
+            square_slice_3d_1d_kernel, dim=a.shape[1:], inputs=[a, b, 1], adjoint=True, adj_inputs=[None, None, 1]
+        )
+
+        expected_grad = [
+            [
+                [0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+            ],
+            [
+                [11 * 2, 12 * 2, 13 * 2],
+                [14 * 2, 15 * 2, 16 * 2],
+                [17 * 2, 18 * 2, 19 * 2],
+            ],
+            [
+                [0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+            ],
+        ]
+        assert_np_equal(a.grad.numpy(), np.array(expected_grad))
+
+
+def test_gradient_slice_3d_2d(test, device):
+    with wp.ScopedDevice(device):
+        data = [
+            [
+                [1, 2, 3],
+                [4, 5, 6],
+                [7, 8, 9],
+            ],
+            [
+                [11, 12, 13],
+                [14, 15, 16],
+                [17, 18, 19],
+            ],
+            [
+                [21, 22, 23],
+                [24, 25, 26],
+                [27, 28, 29],
+            ],
+        ]
+        a = wp.array(data, dtype=float, requires_grad=True)
+        b = wp.zeros_like(a, requires_grad=False)
+        b.grad = wp.ones_like(a, requires_grad=False)
+
+        wp.launch(square_slice_3d_2d_kernel, dim=a.shape[2], inputs=[a, b, 1, 1])
+
+        # use internal gradients (.grad), adj_inputs are None
+        wp.launch(
+            square_slice_3d_2d_kernel, dim=a.shape[2], inputs=[a, b, 1, 1], adjoint=True, adj_inputs=[None, None, 1, 1]
+        )
+
+        expected_grad = [
+            [
+                [0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+            ],
+            [
+                [0, 0, 0],
+                [14 * 2, 15 * 2, 16 * 2],
+                [0, 0, 0],
+            ],
+            [
+                [0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+            ],
+        ]
+        assert_np_equal(a.grad.numpy(), np.array(expected_grad))
+
+
 devices = get_test_devices()
 
 
@@ -739,6 +875,9 @@ add_function_test(TestGrad, "test_aliasing", test_aliasing, devices=devices)
 add_function_test(TestGrad, "test_gradient_internal", test_gradient_internal, devices=devices)
 add_function_test(TestGrad, "test_gradient_external", test_gradient_external, devices=devices)
 add_function_test(TestGrad, "test_gradient_precedence", test_gradient_precedence, devices=devices)
+add_function_test(TestGrad, "test_gradient_slice_2d", test_gradient_slice_2d, devices=devices)
+add_function_test(TestGrad, "test_gradient_slice_3d_1d", test_gradient_slice_3d_1d, devices=devices)
+add_function_test(TestGrad, "test_gradient_slice_3d_2d", test_gradient_slice_3d_2d, devices=devices)
 
 
 if __name__ == "__main__":

warp/tests/test_grad_customs.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 # atomic add function that memorizes which thread incremented the counter
 # so that the correct counter value per thread can be used in the replay

warp/tests/test_hash_grid.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 num_points = 4096
 dim_x = 128
 dim_y = 128
@@ -128,7 +126,7 @@ def test_hashgrid_query(test, device):
 
             print(f"Passed: {np.array_equal(counts, counts_ref)}")
 
-        test.assertTrue(np.array_equal(counts, counts_ref))
+        assert_np_equal(counts, counts_ref)
 
 
 def test_hashgrid_inputs(test, device):