warp-lang 1.6.1__py3-none-macosx_10_13_universal2.whl → 1.7.0__py3-none-macosx_10_13_universal2.whl

warp/examples/fem/example_distortion_energy.py

@@ -1,9 +1,17 @@
-# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# SPDX-FileCopyrightText: Copyright (c) 2022 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 Distortion Energy
@@ -74,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

@@ -1,9 +1,17 @@
-# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# SPDX-FileCopyrightText: Copyright (c) 2024 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 Magnetostatics
@@ -52,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)
 
@@ -77,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
 
@@ -99,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/example_mixed_elasticity.py

@@ -1,9 +1,17 @@
-# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# SPDX-FileCopyrightText: Copyright (c) 2022 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 Mixed Elasticity

warp/examples/fem/example_navier_stokes.py

@@ -1,9 +1,17 @@
-# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# SPDX-FileCopyrightText: Copyright (c) 2022 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 Navier Stokes

warp/examples/fem/example_nonconforming_contact.py

@@ -1,9 +1,17 @@
-# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# SPDX-FileCopyrightText: Copyright (c) 2024 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 Nonconforming Contact

warp/examples/fem/example_stokes.py

@@ -1,9 +1,17 @@
-# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# SPDX-FileCopyrightText: Copyright (c) 2022 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 Stokes

warp/examples/fem/example_stokes_transfer.py

@@ -1,9 +1,17 @@
-# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# SPDX-FileCopyrightText: Copyright (c) 2022 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 Stokes Transfer

warp/examples/fem/example_streamlines.py

@@ -1,9 +1,17 @@
-# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# SPDX-FileCopyrightText: Copyright (c) 2024 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 Streamlines

warp/examples/fem/utils.py

@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
 from typing import Any, Dict, Optional, Tuple
 
 import numpy as np
@@ -19,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
@@ -210,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
 
@@ -230,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)
@@ -443,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),
@@ -485,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()

warp/examples/interop/example_jax_ffi_callback.py

@@ -0,0 +1,132 @@
+# 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 register_ffi_callback()
+#
+# Examples of calling Python functions from JAX.
+# Target functions must have the form func(inputs, outputs, attrs, ctx).
+###########################################################################
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+import warp as wp
+from warp.jax import get_jax_device
+from warp.jax_experimental.ffi import register_ffi_callback
+
+
+@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
+
+
+def example1():
+    # the Python function to call
+    def print_args(inputs, outputs, attrs, ctx):
+        def buffer_to_string(b):
+            return str(b.dtype) + str(list(b.shape)) + " @%x" % b.data
+
+        print("Inputs:     ", ", ".join([buffer_to_string(b) for b in inputs]))
+        print("Outputs:    ", ", ".join([buffer_to_string(b) for b in outputs]))
+        print("Attributes: ", "".join(["\n  %s: %s" % (k, str(v)) for k, v in attrs.items()]))
+
+    # register callback
+    register_ffi_callback("print_args", print_args)
+
+    # set up call
+    call = jax.ffi.ffi_call("print_args", jax.ShapeDtypeStruct((1, 2, 3), jnp.int8))
+
+    # call it
+    call(
+        jnp.arange(16),
+        jnp.arange(32.0).reshape((4, 8)),
+        str_attr="hi",
+        f32_attr=np.float32(4.2),
+        dict_attr={"a": 1, "b": 6.4},
+    )
+
+
+def example2():
+    # the Python function to call
+    def warp_func(inputs, outputs, attrs, ctx):
+        # input arrays
+        a = inputs[0]
+        b = inputs[1]
+
+        # scalar attributes
+        s = attrs["scale"]
+
+        # output arrays
+        c = outputs[0]
+        d = outputs[1]
+
+        device = wp.device_from_jax(get_jax_device())
+        stream = wp.Stream(device, cuda_stream=ctx.stream)
+
+        with wp.ScopedStream(stream):
+            # launch with arrays of scalars
+            wp.launch(scale_kernel, dim=a.shape, inputs=[a, s], outputs=[c])
+
+            # launch with arrays of vec2
+            # NOTE: the input shapes are from JAX arrays, we need to strip the inner dimension for vec2 arrays
+            wp.launch(scale_vec_kernel, dim=b.shape[0], inputs=[b, s], outputs=[d])
+
+    # register callback
+    register_ffi_callback("warp_func", warp_func)
+
+    n = 10
+
+    # inputs
+    a = jnp.arange(n, dtype=jnp.float32)
+    b = jnp.arange(n, dtype=jnp.float32).reshape((n // 2, 2))  # array of wp.vec2
+    s = 2.0
+
+    # set up call
+    out_types = [
+        jax.ShapeDtypeStruct(a.shape, jnp.float32),
+        jax.ShapeDtypeStruct(b.shape, jnp.float32),  # array of wp.vec2
+    ]
+    call = jax.ffi.ffi_call("warp_func", out_types)
+
+    # call it
+    c, d = call(a, b, scale=s)
+
+    print(c)
+    print(d)
+
+
+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()