warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/jax.py

@@ -0,0 +1,29 @@
+# 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.
+
+# isort: skip_file
+
+from warp._src.jax import get_jax_device as get_jax_device
+
+
+# TODO: Remove after cleaning up the public API.
+
+from warp._src import jax as _jax
+
+
+def __getattr__(name):
+    from warp._src.utils import get_deprecated_api
+
+    return get_deprecated_api(_jax, "wp", name)

warp/jax_experimental/__init__.py

@@ -0,0 +1,29 @@
+# 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.
+
+# isort: skip_file
+
+from warp._src.jax_experimental.ffi import GraphMode as GraphMode
+from warp._src.jax_experimental.ffi import jax_kernel as jax_kernel
+from warp._src.jax_experimental.ffi import jax_callable as jax_callable
+from warp._src.jax_experimental.ffi import register_ffi_callback as register_ffi_callback
+
+from warp._src.jax_experimental.ffi import (
+    get_jax_callable_default_graph_cache_max as get_jax_callable_default_graph_cache_max,
+)
+from warp._src.jax_experimental.ffi import (
+    set_jax_callable_default_graph_cache_max as set_jax_callable_default_graph_cache_max,
+)
+from warp._src.jax_experimental.ffi import clear_jax_callable_graph_cache as clear_jax_callable_graph_cache

warp/jax_experimental/custom_call.py

@@ -0,0 +1,29 @@
+# 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.
+
+# isort: skip_file
+
+from warp._src.jax_experimental.custom_call import jax_kernel as jax_kernel
+
+
+# TODO: Remove after cleaning up the public API.
+
+from warp._src.jax_experimental import custom_call as _custom_call
+
+
+def __getattr__(name):
+    from warp._src.utils import get_deprecated_api
+
+    return get_deprecated_api(_custom_call, "wp.jax_experimental", name)

warp/jax_experimental/ffi.py

@@ -0,0 +1,39 @@
+# 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.
+
+# isort: skip_file
+
+from warp._src.jax_experimental.ffi import GraphMode as GraphMode
+from warp._src.jax_experimental.ffi import jax_kernel as jax_kernel
+from warp._src.jax_experimental.ffi import jax_callable as jax_callable
+from warp._src.jax_experimental.ffi import register_ffi_callback as register_ffi_callback
+
+from warp._src.jax_experimental.ffi import (
+    get_jax_callable_default_graph_cache_max as get_jax_callable_default_graph_cache_max,
+)
+from warp._src.jax_experimental.ffi import (
+    set_jax_callable_default_graph_cache_max as set_jax_callable_default_graph_cache_max,
+)
+from warp._src.jax_experimental.ffi import clear_jax_callable_graph_cache as clear_jax_callable_graph_cache
+
+# TODO: Remove after cleaning up the public API.
+
+from warp._src.jax_experimental import ffi as _ffi
+
+
+def __getattr__(name):
+    from warp._src.utils import get_deprecated_api
+
+    return get_deprecated_api(_ffi, "wp.jax_experimental", name)

warp/jax_experimental/xla_ffi.py

@@ -0,0 +1,24 @@
+# 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.
+
+# TODO: Remove after cleaning up the public API.
+
+from warp._src.jax_experimental import xla_ffi as _xla_ffi
+
+
+def __getattr__(name):
+    from warp._src.utils import get_deprecated_api
+
+    return get_deprecated_api(_xla_ffi, "wp.jax_experimental", name)

warp/marching_cubes.py

@@ -0,0 +1,24 @@
+# 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.
+
+# TODO: Remove after cleaning up the public API.
+
+from warp._src import marching_cubes as _marching_cubes
+
+
+def __getattr__(name):
+    from warp._src.utils import get_deprecated_api
+
+    return get_deprecated_api(_marching_cubes, "wp", name)

warp/math.py

@@ -0,0 +1,37 @@
+# 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.
+
+# isort: skip_file
+
+from warp._src.math import norm_huber as norm_huber
+from warp._src.math import norm_l1 as norm_l1
+from warp._src.math import norm_l2 as norm_l2
+from warp._src.math import norm_pseudo_huber as norm_pseudo_huber
+from warp._src.math import smooth_normalize as smooth_normalize
+from warp._src.math import transform_compose as transform_compose
+from warp._src.math import transform_decompose as transform_decompose
+from warp._src.math import transform_from_matrix as transform_from_matrix
+from warp._src.math import transform_to_matrix as transform_to_matrix
+
+
+# TODO: Remove after cleaning up the public API.
+
+from warp._src import math as _math
+
+
+def __getattr__(name):
+    from warp._src.utils import get_deprecated_api
+
+    return get_deprecated_api(_math, "wp", name)