warp-lang 1.7.2rc1__py3-none-win_amd64.whl → 1.8.1__py3-none-win_amd64.whl

warp/tests/test_codegen.py

@@ -435,14 +435,8 @@ def test_error_collection_construct(test, device):
         x = [1.0, 2.0, 3.0]
 
     def kernel_2_fn():
-        x = (1.0, 2.0, 3.0)
-
-    def kernel_3_fn():
         x = {"a": 1.0, "b": 2.0, "c": 3.0}
 
-    def kernel_4_fn():
-        wp.length((1.0, 2.0, 3.0))
-
     kernel = wp.Kernel(func=kernel_1_fn)
     with test.assertRaisesRegex(
         RuntimeError,
@@ -451,22 +445,9 @@ def test_error_collection_construct(test, device):
         wp.launch(kernel, dim=1, device=device)
 
     kernel = wp.Kernel(func=kernel_2_fn)
-    with test.assertRaisesRegex(
-        RuntimeError,
-        r"Tuple constructs are not supported in kernels. Use vectors like `wp.vec3\(\)` for small collections instead.",
-    ):
-        wp.launch(kernel, dim=1, device=device)
-
-    kernel = wp.Kernel(func=kernel_3_fn)
     with test.assertRaisesRegex(RuntimeError, r"Construct `ast.Dict` not supported in kernels."):
         wp.launch(kernel, dim=1, device=device)
 
-    kernel = wp.Kernel(func=kernel_4_fn)
-    with test.assertRaisesRegex(
-        RuntimeError, r"Tuple constructs are not supported in kernels. Use vectors like `wp.vec3\(\)` instead."
-    ):
-        wp.launch(kernel, dim=1, device=device)
-
 
 def test_error_unmatched_arguments(test, device):
     def kernel_1_fn():
@@ -693,6 +674,138 @@ def test_codegen_return_in_kernel(test, device):
     test.assertEqual(result.numpy()[0], grid_size - 256)
 
 
+@wp.kernel
+def conditional_ifexp(x: float, result: wp.array(dtype=wp.int32)):
+    wp.atomic_add(result, 0, 1) if x > 0.0 else wp.atomic_add(result, 1, 1)
+
+
+def test_ifexp_only_executes_one_branch(test, device):
+    result = wp.zeros(2, dtype=wp.int32, device=device)
+
+    wp.launch(conditional_ifexp, dim=1, inputs=[1.0, result], device=device)
+
+    values = result.numpy()
+    # Only first branch is taken
+    test.assertEqual(values[0], 1)
+    test.assertEqual(values[1], 0)
+
+
+@wp.kernel
+def test_multiple_return_values_quat_to_axis_angle_kernel(
+    q: wp.quath,
+    expected_axis: wp.vec3h,
+    expected_angle: wp.float16,
+):
+    axis, angle = wp.quat_to_axis_angle(q)
+
+    wp.expect_near(axis[0], expected_axis[0], tolerance=wp.float16(1e-3))
+    wp.expect_near(axis[1], expected_axis[1], tolerance=wp.float16(1e-3))
+    wp.expect_near(axis[2], expected_axis[2], tolerance=wp.float16(1e-3))
+
+    wp.expect_near(angle, expected_angle, tolerance=wp.float16(1e-3))
+
+
+@wp.kernel
+def test_multiple_return_values_svd3_kernel(
+    A: wp.mat33f,
+    expected_U: wp.mat33f,
+    expected_sigma: wp.vec3f,
+    expected_V: wp.mat33f,
+):
+    U, sigma, V = wp.svd3(A)
+
+    wp.expect_near(U[0][0], expected_U[0][0], tolerance=1e-5)
+    wp.expect_near(U[0][1], expected_U[0][1], tolerance=1e-5)
+    wp.expect_near(U[0][2], expected_U[0][2], tolerance=1e-5)
+    wp.expect_near(U[1][0], expected_U[1][0], tolerance=1e-5)
+    wp.expect_near(U[1][1], expected_U[1][1], tolerance=1e-5)
+    wp.expect_near(U[1][2], expected_U[1][2], tolerance=1e-5)
+    wp.expect_near(U[2][0], expected_U[2][0], tolerance=1e-5)
+    wp.expect_near(U[2][1], expected_U[2][1], tolerance=1e-5)
+    wp.expect_near(U[2][2], expected_U[2][2], tolerance=1e-5)
+
+    wp.expect_near(sigma[0], expected_sigma[0], tolerance=1e-5)
+    wp.expect_near(sigma[1], expected_sigma[1], tolerance=1e-5)
+    wp.expect_near(sigma[2], expected_sigma[2], tolerance=1e-5)
+
+    wp.expect_near(V[0][0], expected_V[0][0], tolerance=1e-5)
+    wp.expect_near(V[0][1], expected_V[0][1], tolerance=1e-5)
+    wp.expect_near(V[0][2], expected_V[0][2], tolerance=1e-5)
+    wp.expect_near(V[1][0], expected_V[1][0], tolerance=1e-5)
+    wp.expect_near(V[1][1], expected_V[1][1], tolerance=1e-5)
+    wp.expect_near(V[1][2], expected_V[1][2], tolerance=1e-5)
+    wp.expect_near(V[2][0], expected_V[2][0], tolerance=1e-5)
+    wp.expect_near(V[2][1], expected_V[2][1], tolerance=1e-5)
+    wp.expect_near(V[2][2], expected_V[2][2], tolerance=1e-5)
+
+
+def test_multiple_return_values(test, device):
+    q = wp.quath(1.0, 2.0, 3.0, 4.0)
+    expected_axis = wp.vec3h(0.26726124, 0.53452247, 0.80178368)
+    expected_angle = 1.50408018
+
+    axis, angle = wp.quat_to_axis_angle(q)
+
+    test.assertAlmostEqual(axis[0], expected_axis[0], places=3)
+    test.assertAlmostEqual(axis[1], expected_axis[1], places=3)
+    test.assertAlmostEqual(axis[2], expected_axis[2], places=3)
+
+    test.assertAlmostEqual(angle, expected_angle, places=3)
+
+    wp.launch(
+        test_multiple_return_values_quat_to_axis_angle_kernel,
+        dim=1,
+        inputs=(q, expected_axis, expected_angle),
+    )
+
+    # fmt: off
+    A = wp.mat33(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0)
+    expected_U = wp.mat33(
+        0.21483721, 0.88723058, -0.40824816,
+        0.52058744, 0.24964368, 0.81649637,
+        0.82633746, -0.38794267, -0.40824834,
+    )
+    expected_sigma = wp.vec3(16.84809875, 1.06836915, 0.00000019)
+    expected_V = wp.mat33(
+        0.47967088, -0.77669406, 0.40824246,
+        0.57236743, -0.07568054, -0.81649727,
+        0.66506463, 0.62531471, 0.40825251,
+    )
+    # fmt: on
+
+    U, sigma, V = wp.svd3(A)
+
+    test.assertAlmostEqual(U[0][0], expected_U[0][0], places=5)
+    test.assertAlmostEqual(U[0][1], expected_U[0][1], places=5)
+    test.assertAlmostEqual(U[0][2], expected_U[0][2], places=5)
+    test.assertAlmostEqual(U[1][0], expected_U[1][0], places=5)
+    test.assertAlmostEqual(U[1][1], expected_U[1][1], places=5)
+    test.assertAlmostEqual(U[1][2], expected_U[1][2], places=5)
+    test.assertAlmostEqual(U[2][0], expected_U[2][0], places=5)
+    test.assertAlmostEqual(U[2][1], expected_U[2][1], places=5)
+    test.assertAlmostEqual(U[2][2], expected_U[2][2], places=5)
+
+    test.assertAlmostEqual(sigma[0], expected_sigma[0], places=5)
+    test.assertAlmostEqual(sigma[1], expected_sigma[1], places=5)
+    test.assertAlmostEqual(sigma[2], expected_sigma[2], places=5)
+
+    test.assertAlmostEqual(V[0][0], expected_V[0][0], places=5)
+    test.assertAlmostEqual(V[0][1], expected_V[0][1], places=5)
+    test.assertAlmostEqual(V[0][2], expected_V[0][2], places=5)
+    test.assertAlmostEqual(V[1][0], expected_V[1][0], places=5)
+    test.assertAlmostEqual(V[1][1], expected_V[1][1], places=5)
+    test.assertAlmostEqual(V[1][2], expected_V[1][2], places=5)
+    test.assertAlmostEqual(V[2][0], expected_V[2][0], places=5)
+    test.assertAlmostEqual(V[2][1], expected_V[2][1], places=5)
+    test.assertAlmostEqual(V[2][2], expected_V[2][2], places=5)
+
+    wp.launch(
+        test_multiple_return_values_svd3_kernel,
+        dim=1,
+        inputs=(A, expected_U, expected_sigma, expected_V),
+    )
+
+
 class TestCodeGen(unittest.TestCase):
     pass
 
@@ -825,6 +938,16 @@ add_kernel_test(TestCodeGen, name="test_call_syntax", kernel=test_call_syntax, d
 add_kernel_test(TestCodeGen, name="test_shadow_builtin", kernel=test_shadow_builtin, dim=1, devices=devices)
 add_kernel_test(TestCodeGen, name="test_while_condition_eval", kernel=test_while_condition_eval, dim=1, devices=devices)
 add_function_test(TestCodeGen, "test_codegen_return_in_kernel", test_codegen_return_in_kernel, devices=devices)
+add_function_test(
+    TestCodeGen, "test_ifexp_only_executes_one_branch", test_ifexp_only_executes_one_branch, devices=devices
+)
+add_function_test(
+    TestCodeGen,
+    func=test_multiple_return_values,
+    name="test_multiple_return_values",
+    devices=devices,
+)
+
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/test_conditional.py

@@ -58,6 +58,48 @@ def test_conditional_if_else_nested():
     wp.expect_eq(e, -2.0)
 
 
+@wp.kernel
+def test_conditional_ifexp():
+    a = 0.5
+    b = 2.0
+
+    c = 1.0 if a > b else -1.0
+
+    wp.expect_eq(c, -1.0)
+
+
+@wp.kernel
+def test_conditional_ifexp_nested():
+    a = 1.0
+    b = 2.0
+
+    c = 3.0 if a > b else 6.0
+    d = 4.0 if a > b else 7.0
+    e = 1.0 if (a > b and c > d) else (-1.0 if a > b else (2.0 if c > d else -2.0))
+
+    wp.expect_eq(e, -2.0)
+
+
+@wp.kernel
+def test_conditional_ifexp_constant():
+    a = 1.0 if False else -1.0
+    b = 2.0 if 123 else -2.0
+
+    wp.expect_eq(a, -1.0)
+    wp.expect_eq(b, 2.0)
+
+
+@wp.kernel
+def test_conditional_ifexp_constant_nested():
+    a = 1.0 if False else (2.0 if True else 3.0)
+    b = 4.0 if 0 else (5.0 if 0 else (6.0 if False else 7.0))
+    c = 8.0 if False else (9.0 if False else (10.0 if 321 else 11.0))
+
+    wp.expect_eq(a, 2.0)
+    wp.expect_eq(b, 7.0)
+    wp.expect_eq(c, 10.0)
+
+
 @wp.kernel
 def test_boolean_and():
     a = 1.0
@@ -90,7 +132,7 @@ def test_boolean_compound():
 
     d = 1.0
 
-    if a > 0.0 and b > 0.0 or c > a:
+    if (a > 0.0 and b > 0.0) or c > a:
         d = -1.0
 
     wp.expect_eq(d, -1.0)
@@ -231,6 +273,10 @@ class TestConditional(unittest.TestCase):
 
 add_kernel_test(TestConditional, kernel=test_conditional_if_else, dim=1, devices=devices)
 add_kernel_test(TestConditional, kernel=test_conditional_if_else_nested, dim=1, devices=devices)
+add_kernel_test(TestConditional, kernel=test_conditional_ifexp, dim=1, devices=devices)
+add_kernel_test(TestConditional, kernel=test_conditional_ifexp_nested, dim=1, devices=devices)
+add_kernel_test(TestConditional, kernel=test_conditional_ifexp_constant, dim=1, devices=devices)
+add_kernel_test(TestConditional, kernel=test_conditional_ifexp_constant_nested, dim=1, devices=devices)
 add_kernel_test(TestConditional, kernel=test_boolean_and, dim=1, devices=devices)
 add_kernel_test(TestConditional, kernel=test_boolean_or, dim=1, devices=devices)
 add_kernel_test(TestConditional, kernel=test_boolean_compound, dim=1, devices=devices)

warp/tests/test_ctypes.py

@@ -541,25 +541,6 @@ def test_scalar_array_types(test, device, load, store):
         )
 
 
-@wp.kernel
-def test_transform_matrix():
-    r = wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.5)
-    t = wp.vec3(0.25, 0.5, -0.75)
-    s = wp.vec3(2.0, 0.5, 0.75)
-
-    m = wp.mat44(t, r, s)
-
-    p = wp.vec3(1.0, 2.0, 3.0)
-
-    r_0 = wp.quat_rotate(r, wp.cw_mul(s, p)) + t
-    r_1 = wp.transform_point(m, p)
-
-    r_2 = wp.transform_vector(m, p)
-
-    wp.expect_near(r_0, r_1, 1.0e-4)
-    wp.expect_near(r_2, r_0 - t, 1.0e-4)
-
-
 devices = get_test_devices()
 
 
@@ -628,7 +609,6 @@ add_function_test(TestCTypes, "test_vec2_transform", test_vec2_transform, device
 add_function_test(TestCTypes, "test_vec3_arg", test_vec3_arg, devices=devices, n=8)
 add_function_test(TestCTypes, "test_vec3_transform", test_vec3_transform, devices=devices, n=8)
 add_function_test(TestCTypes, "test_transform_multiply", test_transform_multiply, devices=devices, n=8)
-add_kernel_test(TestCTypes, name="test_transform_matrix", kernel=test_transform_matrix, dim=1, devices=devices)
 add_function_test(TestCTypes, "test_scalar_array", test_scalar_array, devices=devices)
 add_function_test(TestCTypes, "test_vector_array", test_vector_array, devices=devices)
 

warp/tests/test_devices.py

@@ -70,6 +70,13 @@ def test_devices_can_access_self(test, device):
             test.assertNotEqual(device, device_str)
 
 
+def test_devices_sm_count(test, device):
+    if device.is_cuda:
+        test.assertTrue(device.sm_count > 0)
+    else:
+        test.assertEqual(device.sm_count, 0)
+
+
 devices = get_test_devices()
 
 
@@ -90,6 +97,7 @@ add_function_test(
 )
 add_function_test(TestDevices, "test_devices_verify_cuda_device", test_devices_verify_cuda_device, devices=devices)
 add_function_test(TestDevices, "test_devices_can_access_self", test_devices_can_access_self, devices=devices)
+add_function_test(TestDevices, "test_devices_sm_count", test_devices_sm_count, devices=devices)
 
 
 if __name__ == "__main__":

warp/tests/test_fabricarray.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import math
 import unittest
 from typing import Any
@@ -70,7 +72,7 @@ def _warp_type_to_fabric(dtype, is_array=False):
 
 
 # returns a fabric array interface constructed from a regular array
-def _create_fabric_array_interface(data: wp.array, attrib: str, bucket_sizes: list = None, copy=False):
+def _create_fabric_array_interface(data: wp.array, attrib: str, bucket_sizes: list[int] | None = None, copy=False):
     assert isinstance(data, wp.array)
     assert data.ndim == 1
 
@@ -138,7 +140,7 @@ def _create_fabric_array_interface(data: wp.array, attrib: str, bucket_sizes: li
 
 
 # returns a fabric array array interface constructed from a list of regular arrays
-def _create_fabric_array_array_interface(data: list, attrib: str, bucket_sizes: list = None):
+def _create_fabric_array_array_interface(data: list, attrib: str, bucket_sizes: list[int] | None = None):
     # data should be a list of arrays
     assert isinstance(data, list)
 

warp/tests/test_fem.py

@@ -476,11 +476,8 @@ def _test_geo_cells(
     wp.atomic_add(cell_measures, s.element_index, fem.measure(domain, s) * s.qp_weight)
 
 
-@fem.integrand(kernel_options={"enable_backward": False})
-def _test_cell_lookup(
-    s: fem.Sample,
-    domain: fem.Domain,
-):
+@fem.integrand(kernel_options={"enable_backward": False, "max_unroll": 2})
+def _test_cell_lookup(s: fem.Sample, domain: fem.Domain, cell_filter: wp.array(dtype=int)):
     pos = domain(s)
 
     s_guess = fem.lookup(domain, pos, s)
@@ -491,6 +488,23 @@ def _test_cell_lookup(
     wp.expect_eq(s_noguess.element_index, s.element_index)
     wp.expect_near(domain(s_noguess), pos, 0.001)
 
+    # Filtered lookup
+    max_dist = 10.0
+    filter_target = 1
+    s_filter = fem.lookup(domain, pos, max_dist, cell_filter, filter_target)
+    wp.expect_eq(s_filter.element_index, 0)
+
+    if s.element_index != 0:
+        # test closest point optimality
+        pos_f = domain(s_filter)
+        pos_f += 0.1 * (pos - pos_f)
+        coord_proj, _sq_dist = fem.element_closest_point(domain, s_filter.element_index, pos_f)
+        wp.expect_near(coord_proj, s_filter.element_coords, 0.001)
+
+    # test that extrapolated coordinates yield bak correct position
+    s_filter.element_coords = fem.element_coordinates(domain, s_filter.element_index, pos)
+    wp.expect_near(domain(s_filter), pos, 0.001)
+
 
 @fem.integrand(kernel_options={"enable_backward": False, "max_unroll": 1})
 def _test_geo_sides(
@@ -520,6 +534,7 @@ def _test_geo_sides(
 
     wp.expect_near(coords, inner_side_s.element_coords, 0.0001)
     wp.expect_near(coords, outer_side_s.element_coords, 0.0001)
+    wp.expect_near(coords, fem.element_coordinates(domain, side_index, pos_side), 0.001)
 
     area = fem.measure(domain, s)
     wp.atomic_add(side_measures, side_index, area * s.qp_weight)
@@ -544,7 +559,7 @@ def _test_side_normals(
         wp.expect_near(F_cross[k], nor[k], 0.0001)
 
 
-def _launch_test_geometry_kernel(geo: fem.Geometry, device, test_cell_lookup: bool = True):
+def _launch_test_geometry_kernel(geo: fem.Geometry, device):
     cell_measures = wp.zeros(dtype=float, device=device, shape=geo.cell_count())
     cell_quadrature = fem.RegularQuadrature(fem.Cells(geo), order=2)
 
@@ -557,11 +572,11 @@ def _launch_test_geometry_kernel(geo: fem.Geometry, device, test_cell_lookup: bo
             quadrature=cell_quadrature,
             values={"cell_measures": cell_measures},
         )
-        if test_cell_lookup:
-            fem.interpolate(
-                _test_cell_lookup,
-                quadrature=cell_quadrature,
-            )
+
+        cell_filter = np.zeros(geo.cell_count(), dtype=int)
+        cell_filter[0] = 1
+        cell_filter = wp.array(cell_filter, dtype=int)
+        fem.interpolate(_test_cell_lookup, quadrature=cell_quadrature, values={"cell_filter": cell_filter})
 
         fem.interpolate(
             _test_geo_sides,
@@ -637,14 +652,14 @@ def test_quad_mesh(test, device):
     with wp.ScopedDevice(device):
         positions, quad_vidx = _gen_quadmesh(N)
 
-    geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
+    geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions, build_bvh=True)
 
     test.assertEqual(geo.cell_count(), N**2)
     test.assertEqual(geo.vertex_count(), (N + 1) ** 2)
     test.assertEqual(geo.side_count(), 2 * (N + 1) * N)
     test.assertEqual(geo.boundary_side_count(), 4 * N)
 
-    side_measures, cell_measures = _launch_test_geometry_kernel(geo, device, test_cell_lookup=False)
+    side_measures, cell_measures = _launch_test_geometry_kernel(geo, device)
 
     assert_np_equal(side_measures.numpy(), np.full(side_measures.shape, 1.0 / (N)), tol=1.0e-4)
     assert_np_equal(cell_measures.numpy(), np.full(cell_measures.shape, 1.0 / (N**2)), tol=1.0e-4)
@@ -655,14 +670,14 @@ def test_quad_mesh(test, device):
     positions = np.hstack((positions, np.ones((positions.shape[0], 1))))
     positions = wp.array(positions, device=device, dtype=wp.vec3)
 
-    geo = fem.Quadmesh3D(quad_vertex_indices=quad_vidx, positions=positions)
+    geo = fem.Quadmesh3D(quad_vertex_indices=quad_vidx, positions=positions, build_bvh=True)
 
     test.assertEqual(geo.cell_count(), N**2)
     test.assertEqual(geo.vertex_count(), (N + 1) ** 2)
     test.assertEqual(geo.side_count(), 2 * (N + 1) * N)
     test.assertEqual(geo.boundary_side_count(), 4 * N)
 
-    side_measures, cell_measures = _launch_test_geometry_kernel(geo, device, test_cell_lookup=False)
+    side_measures, cell_measures = _launch_test_geometry_kernel(geo, device)
 
     assert_np_equal(side_measures.numpy(), np.full(side_measures.shape, 1.0 / (N)), tol=1.0e-4)
     assert_np_equal(cell_measures.numpy(), np.full(cell_measures.shape, 1.0 / (N**2)), tol=1.0e-4)
@@ -711,7 +726,7 @@ def test_hex_mesh(test, device):
     with wp.ScopedDevice(device):
         positions, tet_vidx = _gen_hexmesh(N)
 
-    geo = fem.Hexmesh(hex_vertex_indices=tet_vidx, positions=positions)
+    geo = fem.Hexmesh(hex_vertex_indices=tet_vidx, positions=positions, build_bvh=True)
 
     test.assertEqual(geo.cell_count(), (N) ** 3)
     test.assertEqual(geo.vertex_count(), (N + 1) ** 3)
@@ -719,7 +734,7 @@ def test_hex_mesh(test, device):
     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, test_cell_lookup=False)
+    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)
@@ -844,7 +859,8 @@ def test_deformed_geometry(test, device):
         test.assertEqual(geo.side_count(), 6 * (N + 1) * N**2 + (N**3) * 4)
         test.assertEqual(geo.boundary_side_count(), 12 * N * N)
 
-        side_measures, cell_measures = _launch_test_geometry_kernel(deformed_geo, device, test_cell_lookup=False)
+        deformed_geo.build_bvh()
+        side_measures, cell_measures = _launch_test_geometry_kernel(deformed_geo, device)
 
         test.assertAlmostEqual(
             np.sum(cell_measures.numpy()), scale**3, places=4, msg=f"cell_measures = {cell_measures.numpy()}"
@@ -1497,7 +1513,7 @@ def test_point_basis(test, device):
     point_test = fem.make_test(point_space, domain=domain)
 
     # Sample at particle positions
-    ones = fem.integrate(linear_form, fields={"u": point_test}, nodal=True)
+    ones = fem.integrate(linear_form, fields={"u": point_test}, assembly="nodal")
     test.assertAlmostEqual(np.sum(ones.numpy()), 1.0, places=5)
 
     # Sampling outside of particle positions
@@ -1613,12 +1629,23 @@ def test_particle_quadratures(test, device):
     assert_np_equal(measures.grad.numpy(), np.full(3, 4.0))  # == 1.0 / cell_area
 
 
-@fem.integrand
-def _value_at_node(s: fem.Sample, f: fem.Field, values: wp.array(dtype=float)):
+@fem.integrand(kernel_options={"enable_backward": False})
+def _value_at_node(domain: fem.Domain, s: fem.Sample, f: fem.Field, values: wp.array(dtype=float)):
+    # lookup at node is ambiguous, check that partition_lookup retains sample on current partition
+    s_partition = fem.partition_lookup(domain, domain(s))
+    wp.expect_eq(s.element_index, s_partition.element_index)
+    wp.expect_neq(fem.operator.element_partition_index(domain, s.element_index), fem.NULL_ELEMENT_INDEX)
+
     node_index = fem.operator.node_partition_index(f, s.qp_index)
     return values[node_index]
 
 
+@fem.integrand(kernel_options={"enable_backward": False})
+def _test_node_index(s: fem.Sample, u: fem.Field):
+    wp.expect_eq(fem.node_index(u, s), s.qp_index)
+    return 0.0
+
+
 def test_nodal_quadrature(test, device):
     geo = fem.Grid2D(res=wp.vec2i(2))
 
@@ -1635,15 +1662,18 @@ def test_nodal_quadrature(test, device):
 
     # test accessing data associated to a given node
 
-    piecewise_constant_space = fem.make_polynomial_space(geo, degree=0)
-    geo_partition = fem.LinearGeometryPartition(geo, 3, 4)
-    space_partition = fem.make_space_partition(piecewise_constant_space, geo_partition)
+    piecewise_constant_space = fem.make_polynomial_space(geo, degree=1)
+    geo_partition = fem.LinearGeometryPartition(geo, 2, 4)
+    assert geo_partition.cell_count() == 1
+
+    space_partition = fem.make_space_partition(piecewise_constant_space, geo_partition, with_halo=False)
+
     field = fem.make_discrete_field(piecewise_constant_space, space_partition=space_partition)
 
     partition_domain = fem.Cells(geo_partition)
     partition_nodal_quadrature = fem.NodalQuadrature(partition_domain, piecewise_constant_space)
 
-    partition_node_values = wp.array([5.0], dtype=float)
+    partition_node_values = wp.full(value=5.0, shape=space_partition.node_count(), dtype=float)
     val = fem.integrate(
         _value_at_node,
         quadrature=partition_nodal_quadrature,
@@ -1652,6 +1682,9 @@ def test_nodal_quadrature(test, device):
     )
     test.assertAlmostEqual(val, 5.0 / geo.cell_count(), places=5)
 
+    u_test = fem.make_test(space)
+    fem.integrate(_test_node_index, assembly="nodal", fields={"u": u_test})
+
 
 @wp.func
 def aniso_bicubic_fn(x: wp.vec2, scale: wp.vec2):

warp/tests/test_func.py

@@ -259,7 +259,25 @@ def test_return_annotation_none() -> None:
     user_func_return_none()
 
 
-devices = get_test_devices()
+@wp.func
+def divide_by_zero(x: float):
+    return x / 0.0
+
+
+@wp.func
+def normalize_vector(vec_a: wp.vec3):
+    return wp.normalize(vec_a)
+
+
+# This pair is to test the situation where one overload throws an error, but a second one works.
+@wp.func
+def divide_by_zero_overload(x: wp.float32):
+    return x / 0
+
+
+@wp.func
+def divide_by_zero_overload(x: wp.float64):
+    return wp.div(x, 0.0)
 
 
 class TestFunc(unittest.TestCase):
@@ -425,6 +443,29 @@ class TestFunc(unittest.TestCase):
         ):
             a * b
 
+    def test_cpython_call_user_function_with_error(self):
+        # Actually the following also includes a ZeroDivisionError in the message due to exception chaining,
+        # but I don't know how to test for that.
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "Error calling function 'divide_by_zero'. No version succeeded. "
+            "See above for the error from the last version that was tried.",
+        ):
+            divide_by_zero(1.0)
+
+    def test_cpython_call_user_function_with_overloads(self):
+        self.assertEqual(divide_by_zero_overload(1.0), math.inf)
+
+    def test_cpython_call_user_function_with_wrong_argument_types(self):
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "Error calling function 'normalize_vector'. No version succeeded. "
+            "See above for the error from the last version that was tried.",
+        ):
+            normalize_vector(1.0)
+
+
+devices = get_test_devices()
 
 add_kernel_test(TestFunc, kernel=test_overload_func, name="test_overload_func", dim=1, devices=devices)
 add_function_test(TestFunc, func=test_return_func, name="test_return_func", devices=devices)

warp/tests/test_grad.py

@@ -518,7 +518,7 @@ def test_mesh_grad(test, device):
             pos_np[i, j] += eps
             fd_grad[i, j] = (f1 - f2) / (2 * eps)
 
-    assert np.allclose(ad_grad, fd_grad, atol=1e-3)
+    np.testing.assert_allclose(ad_grad, fd_grad, atol=1e-3)
 
 
 @wp.func

warp/tests/test_lerp.py

@@ -119,9 +119,7 @@ def test_lerp(test, device):
                 a = wp.array([test_data.a], dtype=data_type, device=device, requires_grad=True)
                 b = wp.array([test_data.b], dtype=data_type, device=device, requires_grad=True)
                 t = wp.array([test_data.t], dtype=float, device=device, requires_grad=True)
-                out = wp.array(
-                    [0] * wp.types.type_length(data_type), dtype=data_type, device=device, requires_grad=True
-                )
+                out = wp.array([0] * wp.types.type_size(data_type), dtype=data_type, device=device, requires_grad=True)
 
                 with wp.Tape() as tape:
                     wp.launch(kernel, dim=1, inputs=[a, b, t, out], device=device)