warp-lang 1.6.2__py3-none-win_amd64.whl → 1.7.0__py3-none-win_amd64.whl

warp/dlpack.py

@@ -48,10 +48,6 @@ Py_DecRef.restype = None
 
 PyCapsule_Destructor = ctypes.CFUNCTYPE(None, ctypes.c_void_p)
 
-PyCapsule_New = ctypes.pythonapi.PyCapsule_New
-PyCapsule_New.argtypes = [ctypes.c_void_p, ctypes.c_char_p, PyCapsule_Destructor]
-PyCapsule_New.restype = ctypes.py_object
-
 PyCapsule_IsValid = ctypes.pythonapi.PyCapsule_IsValid
 PyCapsule_IsValid.argtypes = [ctypes.py_object, ctypes.c_char_p]
 PyCapsule_IsValid.restype = ctypes.c_int
@@ -105,8 +101,8 @@ def _dlpack_capsule_deleter(ptr) -> None:
 
     capsule = ctypes.cast(ptr, ctypes.py_object)
 
-    if ctypes.pythonapi.PyCapsule_IsValid(capsule, _c_str_dltensor):
-        managed_ptr = ctypes.pythonapi.PyCapsule_GetPointer(capsule, _c_str_dltensor)
+    if PyCapsule_IsValid(capsule, _c_str_dltensor):
+        managed_ptr = PyCapsule_GetPointer(capsule, _c_str_dltensor)
         managed_tensor = DLManagedTensor.from_address(managed_ptr)
         if managed_tensor.deleter:
             managed_tensor.deleter(managed_ptr)
@@ -302,6 +298,12 @@ def to_dlpack(wp_array: warp.array):
 
     managed_tensor.deleter = _dlpack_tensor_deleter
 
+    # NOTE: jax.ffi.pycapsule() defines the PyCapsule_New() argtypes incorrectly, which causes problems.
+    # Here we make sure that the PyCapsule_Destructor callback is correctly defined.
+    PyCapsule_New = ctypes.pythonapi.PyCapsule_New
+    PyCapsule_New.argtypes = [ctypes.c_void_p, ctypes.c_char_p, PyCapsule_Destructor]
+    PyCapsule_New.restype = ctypes.py_object
+
     capsule = PyCapsule_New(
         ctypes.byref(managed_tensor),
         _c_str_dltensor,

warp/examples/benchmarks/benchmark_tile_load_store.py

@@ -0,0 +1,103 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+import warp as wp
+
+BLOCK_DIM = 128
+
+TILE = 32
+
+
+def create_test_kernel(storage_type: str):
+    @wp.kernel
+    def load_store(a: wp.array2d(dtype=wp.float32), b: wp.array2d(dtype=wp.float32)):
+        i, j = wp.tid()
+
+        if wp.static(storage_type == "shared"):
+            a_tile = wp.tile_load(a, shape=(TILE, TILE), offset=(i * TILE, j * TILE), storage="shared")
+        else:
+            a_tile = wp.tile_load(a, shape=(TILE, TILE), offset=(i * TILE, j * TILE), storage="register")
+
+        wp.tile_store(b, a_tile, offset=(i * TILE, j * TILE))
+
+    return load_store
+
+
+if __name__ == "__main__":
+    wp.config.quiet = True
+    wp.init()
+    wp.clear_kernel_cache()
+    wp.set_module_options({"fast_math": True, "enable_backward": False})
+
+    iterations = 100
+    rng = np.random.default_rng(42)
+
+    shared_benchmark_data = {}
+    register_benchmark_data = {}
+    memcpy_benchmark_data = {}
+
+    sizes = list(range(128, 4097, 128))
+
+    print(f"{'Transfer Size (Bytes)':<23s} {'Shared (GiB/s)':<16s} {'Register (GiB/s)':<18s} {'memcpy (GiB/s)':<16s}")
+    print("-" * 79)
+
+    for size in sizes:
+        a = wp.array(rng.random((size, size), dtype=np.float32), dtype=wp.float32)
+        b = wp.empty_like(a)
+
+        for storage_type in ("shared", "register"):
+            load_store = create_test_kernel(storage_type)
+
+            cmd = wp.launch_tiled(
+                load_store,
+                dim=(a.shape[0] // TILE, a.shape[1] // TILE),
+                inputs=[a],
+                outputs=[b],
+                block_dim=BLOCK_DIM,
+                record_cmd=True,
+            )
+            # Warmup
+            for _ in range(5):
+                cmd.launch()
+
+            with wp.ScopedTimer("benchmark", cuda_filter=wp.TIMING_KERNEL, print=False, synchronize=True) as timer:
+                for _ in range(iterations):
+                    cmd.launch()
+
+            np.testing.assert_equal(a.numpy(), b.numpy())
+
+            timing_results = [result.elapsed for result in timer.timing_results]
+            avg_bw = 2.0 * (a.capacity / (1024 * 1024 * 1024)) / (1e-3 * np.mean(timing_results))
+
+            if storage_type == "shared":
+                shared_benchmark_data[a.capacity] = avg_bw
+            else:
+                register_benchmark_data[a.capacity] = avg_bw
+
+        # Compare with memcpy
+        with wp.ScopedTimer("benchmark", cuda_filter=wp.TIMING_MEMCPY, print=False, synchronize=True) as timer:
+            for _ in range(iterations):
+                wp.copy(b, a)
+
+        timing_results = [result.elapsed for result in timer.timing_results]
+        avg_bw = 2.0 * (a.capacity / (1024 * 1024 * 1024)) / (1e-3 * np.mean(timing_results))
+        memcpy_benchmark_data[a.capacity] = avg_bw
+
+        # Print results
+        print(
+            f"{a.capacity:<23d} {shared_benchmark_data[a.capacity]:<#16.4g} {register_benchmark_data[a.capacity]:<#18.4g} {memcpy_benchmark_data[a.capacity]:<#16.4g}"
+        )

warp/examples/core/example_sample_mesh.py

@@ -0,0 +1,300 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Sample Mesh
+#
+# Shows how to sample points on a mesh's surface using
+# a Cumulative Distribution Function (CDF).
+#
+# The CDF enables uniform sampling of points across the mesh's surface,
+# even when the density of triangles varies. It represents the cumulative
+# probability of selecting a triangle from the mesh, with each triangle
+# weighted by its area relative to the total surface area of the mesh.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+import warp.render
+
+# fmt: off
+POINTS = np.array(
+    (
+        (-0.986598, -0.400638, -0.175759), (-0.81036 , -0.482105, -0.541125),
+        (-1.079616,  0.022652, -0.023381), (-0.894468, -0.080795, -0.618379),
+        (-0.607365, -0.702012, -0.556551), (-0.366107, -0.800096, -0.620734),
+        (-0.801777, -0.690991, -0.239593), (-0.553576, -0.871746, -0.335518),
+        (-0.309133, -0.370805, -0.965784), (-0.288299, -0.956987, -0.402091),
+        (-0.051878, -0.894342, -0.597583), (-0.386774, -1.003107, -0.145116),
+        (-0.19062 , -1.061165,  0.012418), (-0.176053, -1.044838, -0.217194),
+        ( 0.001479, -1.020045, -0.356905), (-0.105375, -0.655117, -0.861365),
+        (-0.542102, -0.517255, -0.795259), (-0.476599, -0.105709, -0.981171),
+        (-1.047915, -0.121584,  0.322098), (-0.527852,  0.137252,  0.501813),
+        (-0.721762, -0.803275,  0.117162), (-0.904992, -0.573281,  0.168408),
+        (-0.796762, -0.473428,  0.569649), (-0.606446, -0.753374,  0.492938),
+        (-0.466481, -0.576566,  0.802562), (-0.50476 , -0.908596,  0.300064),
+        (-0.337425, -1.008902,  0.170911), (-0.048676, -1.055594,  0.246732),
+        (-0.212871, -0.760442,  0.738447), (-0.281356, -0.9322  ,  0.474965),
+        (-0.560476,  0.062512, -0.561019), (-0.003252,  0.083237, -1.049784),
+        (-0.009392,  0.593703, -0.522479), (-0.530465,  0.577231,  0.007172),
+        (-0.02106 ,  0.064189,  1.066722), (-0.003512,  0.59714 ,  0.516904),
+        ( 0.000194,  1.093899,  0.001113), ( 0.256861, -0.955856, -0.445325),
+        ( 0.251205, -1.038759, -0.174212), ( 0.170201, -0.800019, -0.712158),
+        ( 0.364385, -0.560298, -0.866843), ( 0.092809, -0.269437, -1.058467),
+        ( 0.628127, -0.12359 , -0.9012  ), ( 0.507433, -0.930658, -0.215908),
+        ( 0.496448, -0.800205, -0.545904), ( 0.757415, -0.527449, -0.565395),
+        ( 0.908704, -0.596257,  0.028995), ( 0.754069, -0.731365, -0.256687),
+        ( 0.921362, -0.09028 , -0.546421), ( 1.017846, -0.335787, -0.263017),
+        ( 0.016768, -1.080014, -0.058473), ( 0.204245, -1.056388,  0.078346),
+        ( 0.260892, -1.001704,  0.322104), ( 0.16608 , -0.739172,  0.788097),
+        ( 0.021091, -0.931327,  0.557789), (-0.046158, -0.408417,  1.011046),
+        ( 0.429623, -0.987237,  0.088537), ( 0.704993, -0.739396,  0.386838),
+        ( 0.37277 , -0.825639,  0.591102), ( 0.493947, -0.896091,  0.339163),
+        ( 0.321112, -0.540547,  0.890161), ( 0.654753, -0.520495,  0.690104),
+        ( 0.922472, -0.124429,  0.530498), ( 0.662544, -0.85601 ,  0.054375),
+        ( 0.950976, -0.422783,  0.327726), ( 0.536849,  0.109943, -0.52279 ),
+        ( 0.517242,  0.120634,  0.535708), ( 0.532707,  0.598943, -0.000767),
+        ( 1.086691,  0.048722,  0.032517), ( 0.528734, -0.109809,  0.96863 ),
+        (-0.581832, -0.916941, -0.027829), (-0.625071, -0.14445 ,  0.906538),
+    ),
+    dtype=np.float32,
+)
+
+FACE_VERTEX_INDICES = np.array(
+    (
+         6,  0,  1,  6, 21,  0,  2,  0, 18,  0,  3,  1,  2,  3,  0,  5,
+         7,  4, 70,  7, 11,  4,  6,  1, 16,  1,  3,  7,  6,  4,  4,  1,
+        16,  9,  7,  5,  3, 17, 16, 16, 17,  8, 41,  8, 17, 30, 17,  3,
+        10, 14,  9,  5, 10,  9, 10, 37, 14, 15, 10,  5,  7,  9, 11, 11,
+         9, 13, 11, 13, 12, 50, 12, 13,  9, 14, 13, 15, 16,  8, 15,  8,
+        41, 16,  5,  4, 16, 15,  5, 17, 31, 41, 21, 22, 18, 20, 21,  6,
+        18,  0, 21, 20, 25, 23, 20, 70, 25, 70, 11, 26, 26, 25, 70, 25,
+        29, 23, 21, 20, 23, 21, 23, 22, 23, 24, 22, 24, 71, 22, 26, 29,
+        25, 26, 11, 12, 12, 27, 26, 26, 27, 29, 27, 54, 29, 27, 12, 50,
+        28, 29, 54, 54, 53, 28, 23, 28, 24, 29, 28, 23, 28, 55, 24, 28,
+        53, 55, 53, 60, 55, 24, 55, 71, 55, 34, 71, 30,  3,  2,  2, 33,
+        30, 17, 30, 31, 32, 31, 30, 33, 36, 32, 19, 33,  2, 19, 35, 33,
+        19, 71, 34, 35, 19, 34, 34, 66, 35, 35, 36, 33, 35, 67, 36, 15,
+        39, 10, 10, 39, 37, 44, 37, 39, 14, 50, 13, 14, 38, 50, 14, 37,
+        38, 37, 43, 38, 40, 15, 41, 40, 39, 15, 41, 42, 40, 44, 39, 40,
+        31, 42, 41, 38, 43, 56, 44, 43, 37, 44, 47, 43, 47, 63, 43, 44,
+        40, 45, 42, 45, 40, 46, 63, 47, 45, 47, 44, 65, 48, 42, 46, 47,
+        49, 49, 47, 45, 48, 45, 42, 45, 48, 49, 68, 49, 48, 27, 52, 54,
+        50, 51, 27, 27, 51, 52, 50, 38, 51, 38, 56, 51, 51, 56, 52, 54,
+        52, 58, 52, 59, 58, 53, 54, 58, 60, 69, 55, 55, 69, 34, 43, 63,
+        56, 59, 52, 56, 63, 59, 56, 63, 57, 59, 58, 60, 53, 57, 58, 59,
+        58, 57, 61, 60, 58, 61, 57, 64, 61, 62, 61, 64, 60, 61, 69, 62,
+        69, 61, 46, 57, 63, 64, 57, 46, 46, 49, 64, 68, 64, 49, 62, 64,
+        68, 32, 65, 31, 65, 32, 67, 32, 36, 67, 65, 42, 31, 67, 68, 65,
+        48, 65, 68, 34, 69, 66, 67, 35, 66, 68, 66, 62, 66, 69, 62, 67,
+        66, 68, 33, 32, 30, 19,  2, 18, 20,  6, 70,  7, 70,  6, 18, 71,
+        19, 22, 71, 18,
+    ),
+    dtype=np.int32,
+)
+# fmt: on
+
+
+@wp.kernel(enable_backward=False)
+def compute_tri_areas(
+    points: wp.array(dtype=wp.vec3),
+    face_vertex_indices: wp.array(dtype=wp.int32),
+    out_tri_areas: wp.array(dtype=wp.float32),
+    out_total_area: wp.array(dtype=wp.float32),
+):
+    tri = wp.tid()
+
+    # Retrieve the indices of the three vertices that form the current triangle.
+    vtx_0 = face_vertex_indices[tri * 3]
+    vtx_1 = face_vertex_indices[tri * 3 + 1]
+    vtx_2 = face_vertex_indices[tri * 3 + 2]
+
+    # Retrieve their 3D position.
+    pt_0 = points[vtx_0]
+    pt_1 = points[vtx_1]
+    pt_2 = points[vtx_2]
+
+    # Calculate the cross product of two edges of the triangle,
+    # which gives a vector whose magnitude is twice the area of the triangle.
+    cross = wp.cross((pt_1 - pt_0), (pt_2 - pt_0))
+    area = wp.length(cross) * 0.5
+
+    # Store the result.
+    out_tri_areas[tri] = area
+    wp.atomic_add(out_total_area, 0, area)
+
+
+@wp.kernel(enable_backward=False)
+def compute_probability_distribution(
+    tri_areas: wp.array(dtype=wp.float32),
+    total_area: wp.array(dtype=wp.float32),
+    out_probabilities: wp.array(dtype=wp.float32),
+):
+    tri = wp.tid()
+
+    # Calculate the probability of selecting this triangle,
+    # which is proportional to the triangle's area relative to total mesh area.
+    out_probabilities[tri] = tri_areas[tri] / total_area[0]
+
+
+@wp.kernel(enable_backward=False)
+def accumulate_cdf(
+    tri_count: wp.int32,
+    out_cdf: wp.array(dtype=wp.float32),
+):
+    # Transform probability values into a Cumulative Distribution Function (CDF).
+    for tri in range(1, tri_count):
+        out_cdf[tri] += out_cdf[tri - 1]
+
+
+@wp.kernel(enable_backward=False)
+def sample_mesh(
+    mesh: wp.uint64,
+    cdf: wp.array(dtype=wp.float32),
+    seed: wp.int32,
+    out_points: wp.array(dtype=wp.vec3),
+):
+    tid = wp.tid()
+
+    rng = wp.rand_init(seed, tid)
+
+    # Sample the triangle index using the CDF.
+    sample = wp.randf(rng)
+    tri = wp.lower_bound(cdf, sample)
+
+    # Sample the location in that triangle using random barycentric cordinates.
+    ru = wp.randf(rng)
+    rv = wp.randf(rng)
+    tri_u = 1.0 - wp.sqrt(ru)
+    tri_v = wp.sqrt(ru) * (1.0 - rv)
+    pos = wp.mesh_eval_position(mesh, tri, tri_u, tri_v)
+
+    # Store the result.
+    out_points[tid] = pos
+
+
+class Example:
+    def __init__(self, stage_path="example_sample_mesh.usd"):
+        self.mesh = wp.Mesh(
+            points=wp.array(POINTS, dtype=wp.vec3),
+            indices=wp.array(FACE_VERTEX_INDICES, dtype=wp.int32),
+        )
+        self.tri_count = len(FACE_VERTEX_INDICES) // 3
+
+        # Compute the area of each triangle and the total area of the mesh.
+        tri_areas = wp.empty(shape=(self.tri_count,), dtype=wp.float32)
+        total_area = wp.zeros(shape=(1,), dtype=wp.float32)
+        wp.launch(
+            compute_tri_areas,
+            dim=tri_areas.shape,
+            inputs=(
+                self.mesh.points,
+                self.mesh.indices,
+            ),
+            outputs=(
+                tri_areas,
+                total_area,
+            ),
+        )
+
+        # Build a Cumulative Distribution Function (CDF) where the probability
+        # of sampling a given triangle is proportional to its area.
+        self.cdf = wp.empty(shape=(self.tri_count,), dtype=wp.float32)
+        wp.launch(
+            compute_probability_distribution,
+            dim=self.cdf.shape,
+            inputs=(
+                tri_areas,
+                total_area,
+            ),
+            outputs=(self.cdf,),
+        )
+        wp.launch(
+            accumulate_cdf,
+            dim=(1,),
+            inputs=(self.tri_count,),
+            outputs=(self.cdf,),
+        )
+
+        # Array to store the sampled points.
+        self.points = wp.empty(shape=(100,), dtype=wp.vec3)
+
+        self.fps = 4
+        self.frame = 0
+
+        if stage_path:
+            self.renderer = wp.render.UsdRenderer(stage_path, fps=self.fps)
+        else:
+            self.renderer = None
+
+    def step(self):
+        with wp.ScopedTimer("step"):
+            # Sample new points on the mesh using the CDF and the current frame
+            # number as seed to ensure different samples each frame.
+            wp.launch(
+                sample_mesh,
+                dim=self.points.shape,
+                inputs=(
+                    self.mesh.id,
+                    self.cdf,
+                    self.frame,
+                ),
+                outputs=(self.points,),
+            )
+
+            self.frame += 1
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        with wp.ScopedTimer("render"):
+            self.renderer.begin_frame(self.frame / self.fps)
+            self.renderer.render_mesh(
+                name="mesh",
+                points=self.mesh.points.numpy(),
+                indices=self.mesh.indices.numpy(),
+                colors=(0.35, 0.55, 0.9),
+            )
+            self.renderer.render_points(name="points", points=self.points.numpy(), radius=0.05, colors=(0.8, 0.3, 0.2))
+            self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument(
+        "--stage_path",
+        type=lambda x: None if x == "None" else str(x),
+        default="example_sample_mesh.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--num_frames", type=int, default=16, help="Total number of frames.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(stage_path=args.stage_path)
+
+        for _ in range(args.num_frames):
+            example.step()
+            example.render()
+
+        if example.renderer:
+            example.renderer.save()

warp/examples/fem/example_apic_fluid.py

@@ -117,7 +117,7 @@ def divergence_form(s: Sample, domain: Domain, u: Field, psi: Field):
 def invert_volume_kernel(values: wp.array(dtype=float)):
     i = wp.tid()
     m = values[i]
-    values[i] = wp.select(m == 0.0, 1.0 / m, 0.0)
+    values[i] = wp.where(m == 0.0, 0.0, 1.0 / m)
 
 
 @wp.kernel

warp/examples/fem/example_burgers.py

@@ -75,7 +75,7 @@ def cell_transport_form(s: fem.Sample, domain: fem.Domain, u: fem.Field, v: fem.
 def initial_condition(s: fem.Sample, domain: fem.Domain):
     x = domain(s)[0] * 2.0
     wave = wp.sin(x * wp.pi)
-    return wp.vec2(wp.select(x <= 1.0, 0.0, wave), 0.0)
+    return wp.vec2(wp.where(x <= 1.0, wave, 0.0), 0.0)
 
 
 @fem.integrand
@@ -87,7 +87,7 @@ def velocity_norm(s: fem.Sample, u: fem.Field):
 def minmod(a: float, b: float):
     sa = wp.sign(a)
     sb = wp.sign(b)
-    return wp.select(sa == sb, 0.0, sa * wp.min(wp.abs(a), wp.abs(b)))
+    return wp.where(sa == sb, sa * wp.min(wp.abs(a), wp.abs(b)), 0.0)
 
 
 @fem.integrand

warp/examples/fem/example_deformed_geometry.py

@@ -57,7 +57,7 @@ def boundary_projector_form(
     Bilinear boundary condition projector form, non-zero on radial edges
     """
     nor = fem.normal(domain, s)
-    active = wp.select(nor[0] < -0.9999 or nor[1] < -0.9999, 0.0, 1.0)
+    active = wp.where(nor[0] < -0.9999 or nor[1] < -0.9999, 1.0, 0.0)
     return active * u(s) * v(s)
 
 

warp/examples/fem/example_distortion_energy.py

@@ -82,7 +82,7 @@ def boundary_projector_form(
 ):
     # Fix a single point
     # (underconstrained, solution up to a rotation in UV space)
-    w = wp.select(s.qp_index == 0, 0.0, 1.0)
+    w = wp.where(s.qp_index == 0, 1.0, 0.0)
     return w * wp.dot(u(s), v(s))
 
 

warp/examples/fem/example_magnetostatics.py

@@ -60,8 +60,8 @@ def cube_to_cylinder_grad(x: wp.vec3):
         dir_grad = (wp.identity(n=3, dtype=float) - wp.outer(dir_xz, dir_xz)) / wp.length(pos_xz)
 
         abs_xz = wp.abs(pos_xz)
-        xinf_grad = wp.select(
-            abs_xz[0] > abs_xz[2], wp.vec3(0.0, 0.0, wp.sign(pos_xz[2])), wp.vec(wp.sign(pos_xz[0]), 0.0, 0.0)
+        xinf_grad = wp.where(
+            abs_xz[0] > abs_xz[2], wp.vec(wp.sign(pos_xz[0]), 0.0, 0.0), wp.vec3(0.0, 0.0, wp.sign(pos_xz[2]))
         )
         grad = dir_grad * wp.max(abs_xz) + wp.outer(dir_xz, xinf_grad)
 
@@ -85,10 +85,10 @@ def permeability_field(
     r = wp.sqrt(x * x + z * z)
 
     if r <= core_radius:
-        return wp.select(y < core_height, MU_0, MU_i)
+        return wp.where(y < core_height, MU_i, MU_0)
 
     if r >= coil_internal_radius and r <= coil_external_radius:
-        return wp.select(y < coil_height, MU_0, MU_c)
+        return wp.where(y < coil_height, MU_c, MU_0)
 
     return MU_0
 
@@ -107,10 +107,10 @@ def current_field(
 
     r = wp.sqrt(x * x + z * z)
 
-    return wp.select(
+    return wp.where(
         y < coil_height and r >= coil_internal_radius and r <= coil_external_radius,
-        wp.vec3(0.0),
         wp.vec3(z, 0.0, -x) * current / r,
+        wp.vec3(0.0),
     )
 
 

warp/examples/fem/utils.py

@@ -34,6 +34,9 @@ __all__ = [
     "Plot",
 ]
 
+# matrix inversion routines contain nested loops,
+# default unrolling leads to code explosion
+wp.set_module_options({"max_unroll": 6})
 
 #
 # Mesh utilities
@@ -225,6 +228,7 @@ def bsr_cg(
     mv_routine=None,
     quiet=False,
     method: str = "cg",
+    M: BsrMatrix = None,
 ) -> Tuple[float, int]:
     """Solves the linear system A x = b using an iterative solver, optionally with diagonal preconditioning
 
@@ -245,7 +249,9 @@ def bsr_cg(
 
     """
 
-    if mv_routine is None:
+    if M is not None:
+        M = aslinearoperator(M)
+    elif mv_routine is None:
         M = preconditioner(A, "diag") if use_diag_precond else None
     else:
         A = LinearOperator(A.shape, A.dtype, A.device, matvec=mv_routine)
@@ -458,7 +464,7 @@ def bsr_solve_saddle(
     return err, end_iter
 
 
-@wp.kernel
+@wp.kernel(enable_backward=False)
 def _compute_schur_inverse_diagonal(
     B_offsets: wp.array(dtype=int),
     B_indices: wp.array(dtype=int),
@@ -500,7 +506,7 @@ def invert_diagonal_bsr_matrix(A: BsrMatrix):
     )
 
 
-@wp.kernel
+@wp.kernel(enable_backward=False)
 def _block_diagonal_invert(values: wp.array(dtype=Any)):
     i = wp.tid()
     values[i] = fem.utils.inverse_qr(values[i])

warp/examples/interop/example_jax_callable.py

@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example jax_callable()
+#
+# Examples of calling annotated Python functions from JAX.
+###########################################################################
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+import warp as wp
+from warp.jax_experimental.ffi import jax_callable
+
+
+@wp.kernel
+def scale_kernel(a: wp.array(dtype=float), s: float, output: wp.array(dtype=float)):
+    tid = wp.tid()
+    output[tid] = a[tid] * s
+
+
+@wp.kernel
+def scale_vec_kernel(a: wp.array(dtype=wp.vec2), s: float, output: wp.array(dtype=wp.vec2)):
+    tid = wp.tid()
+    output[tid] = a[tid] * s
+
+
+# The Python function to call.
+# Note the argument annotations, just like Warp kernels.
+def example_func(
+    # inputs
+    a: wp.array(dtype=float),
+    b: wp.array(dtype=wp.vec2),
+    s: float,
+    # outputs
+    c: wp.array(dtype=float),
+    d: wp.array(dtype=wp.vec2),
+):
+    wp.launch(scale_kernel, dim=a.shape, inputs=[a, s], outputs=[c])
+    wp.launch(scale_vec_kernel, dim=b.shape, inputs=[b, s], outputs=[d])
+
+
+def example1():
+    jax_func = jax_callable(example_func, num_outputs=2, vmap_method="broadcast_all")
+
+    @jax.jit
+    def f():
+        # inputs
+        a = jnp.arange(10, dtype=jnp.float32)
+        b = jnp.arange(10, dtype=jnp.float32).reshape((5, 2))  # wp.vec2
+        s = 2.0
+
+        # output shapes
+        output_dims = {"c": a.shape, "d": b.shape}
+
+        c, d = jax_func(a, b, s, output_dims=output_dims)
+
+        return c, d
+
+    r1, r2 = f()
+    print(r1)
+    print(r2)
+
+
+def example2():
+    jax_func = jax_callable(example_func, num_outputs=2, vmap_method="broadcast_all")
+
+    # NOTE: scalar arguments must be static compile-time constants
+    @partial(jax.jit, static_argnames=["s"])
+    def f(a, b, s):
+        # output shapes
+        output_dims = {"c": a.shape, "d": b.shape}
+
+        c, d = jax_func(a, b, s, output_dims=output_dims)
+
+        return c, d
+
+    # inputs
+    a = jnp.arange(10, dtype=jnp.float32)
+    b = jnp.arange(10, dtype=jnp.float32).reshape((5, 2))  # wp.vec2
+    s = 3.0
+
+    r1, r2 = f(a, b, s)
+    print(r1)
+    print(r2)
+
+
+def main():
+    wp.init()
+    wp.load_module(device=wp.get_device())
+
+    examples = [example1, example2]
+
+    for example in examples:
+        print("\n===========================================================================")
+        print(f"{example.__name__}:")
+        example()
+
+
+if __name__ == "__main__":
+    main()