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

warp/autograd.py

@@ -13,34 +13,36 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import inspect
 import itertools
-from typing import Any, Callable, Dict, List, Sequence, Tuple, Union
+from typing import Any, Callable, Sequence
 
 import numpy as np
 
 import warp as wp
 
 __all__ = [
-    "jacobian",
-    "jacobian_fd",
     "gradcheck",
     "gradcheck_tape",
+    "jacobian",
+    "jacobian_fd",
     "jacobian_plot",
 ]
 
 
 def gradcheck(
-    function: Union[wp.Kernel, Callable],
-    dim: Tuple[int] = None,
-    inputs: Sequence = None,
-    outputs: Sequence = None,
+    function: wp.Kernel | Callable,
+    dim: tuple[int] | None = None,
+    inputs: Sequence | None = None,
+    outputs: Sequence | None = None,
     *,
     eps: float = 1e-4,
     atol: float = 1e-3,
     rtol: float = 1e-2,
     raise_exception: bool = True,
-    input_output_mask: List[Tuple[Union[str, int], Union[str, int]]] = None,
+    input_output_mask: list[tuple[str | int, str | int]] | None = None,
     device: wp.context.Devicelike = None,
     max_blocks: int = 0,
     block_dim: int = 256,
@@ -242,9 +244,9 @@ def gradcheck_tape(
     atol=1e-3,
     rtol=1e-2,
     raise_exception=True,
-    input_output_masks: Dict[str, List[Tuple[Union[str, int], Union[str, int]]]] = None,
-    blacklist_kernels: List[str] = None,
-    whitelist_kernels: List[str] = None,
+    input_output_masks: dict[str, list[tuple[str | int, str | int]]] | None = None,
+    blacklist_kernels: list[str] | None = None,
+    whitelist_kernels: list[str] | None = None,
     max_inputs_per_var=-1,
     max_outputs_per_var=-1,
     plot_relative_error=False,
@@ -364,13 +366,13 @@ class FunctionMetadata:
 
     def __init__(
         self,
-        key: str = None,
-        input_labels: List[str] = None,
-        output_labels: List[str] = None,
-        input_strides: List[tuple] = None,
-        output_strides: List[tuple] = None,
-        input_dtypes: list = None,
-        output_dtypes: list = None,
+        key: str | None = None,
+        input_labels: list[str] | None = None,
+        output_labels: list[str] | None = None,
+        input_strides: list[tuple] | None = None,
+        output_strides: list[tuple] | None = None,
+        input_dtypes: list | None = None,
+        output_dtypes: list | None = None,
     ):
         self.key = key
         self.input_labels = input_labels
@@ -413,7 +415,7 @@ class FunctionMetadata:
                 self.output_strides.append(None)
                 self.output_dtypes.append(None)
 
-    def update_from_function(self, function: Callable, inputs: Sequence, outputs: Sequence = None):
+    def update_from_function(self, function: Callable, inputs: Sequence, outputs: Sequence | None = None):
         self.key = function.__name__
         self.input_labels = list(inspect.signature(function).parameters.keys())
         if outputs is None:
@@ -442,16 +444,15 @@ class FunctionMetadata:
 
 
 def jacobian_plot(
-    jacobians: Dict[Tuple[int, int], wp.array],
-    kernel: Union[FunctionMetadata, wp.Kernel],
-    inputs: Sequence = None,
-    outputs: Sequence = None,
-    show_plot=True,
-    show_colorbar=True,
-    scale_colors_per_submatrix=False,
-    title: str = None,
+    jacobians: dict[tuple[int, int], wp.array],
+    kernel: FunctionMetadata | wp.Kernel,
+    inputs: Sequence | None = None,
+    show_plot: bool = True,
+    show_colorbar: bool = True,
+    scale_colors_per_submatrix: bool = False,
+    title: str | None = None,
     colormap: str = "coolwarm",
-    log_scale=False,
+    log_scale: bool = False,
 ):
     """
     Visualizes the Jacobians computed by :func:`jacobian` or :func:`jacobian_fd` in a combined image plot.
@@ -461,7 +462,6 @@ def jacobian_plot(
         jacobians: A dictionary of Jacobians, where the keys are tuples of input and output indices, and the values are the Jacobian matrices.
         kernel: The Warp kernel function, decorated with the ``@wp.kernel`` decorator, or a :class:`FunctionMetadata` instance with the kernel/function attributes.
         inputs: List of input variables.
-        outputs: List of output variables. Deprecated and will be removed in a future Warp version.
         show_plot: If True, displays the plot via ``plt.show()``.
         show_colorbar: If True, displays a colorbar next to the plot (or a colorbar next to every submatrix if ).
         scale_colors_per_submatrix: If True, considers the minimum and maximum of each Jacobian submatrix separately for color scaling. Otherwise, uses the global minimum and maximum of all Jacobians.
@@ -484,12 +484,6 @@ def jacobian_plot(
         metadata = kernel
     else:
         raise ValueError("Invalid kernel argument: must be a Warp kernel or a FunctionMetadata object")
-    if outputs is not None:
-        wp.utils.warn(
-            "The `outputs` argument to `jacobian_plot` is no longer needed and will be removed in a future Warp version.",
-            DeprecationWarning,
-            stacklevel=3,
-        )
 
     jacobians = sorted(jacobians.items(), key=lambda x: (x[0][1], x[0][0]))
     jacobians = dict(jacobians)
@@ -635,59 +629,6 @@ def jacobian_plot(
     return fig
 
 
-def plot_kernel_jacobians(
-    jacobians: Dict[Tuple[int, int], wp.array],
-    kernel: wp.Kernel,
-    inputs: Sequence,
-    outputs: Sequence,
-    show_plot=True,
-    show_colorbar=True,
-    scale_colors_per_submatrix=False,
-    title: str = None,
-    colormap: str = "coolwarm",
-    log_scale=False,
-):
-    """
-    Visualizes the Jacobians computed by :func:`jacobian` or :func:`jacobian_fd` in a combined image plot.
-    Requires the ``matplotlib`` package to be installed.
-
-    Note:
-        This function is deprecated and will be removed in a future Warp version. Please call :func:`jacobian_plot` instead.
-
-    Args:
-        jacobians: A dictionary of Jacobians, where the keys are tuples of input and output indices, and the values are the Jacobian matrices.
-        kernel: The Warp kernel function, decorated with the ``@wp.kernel`` decorator.
-        inputs: List of input variables.
-        outputs: List of output variables.
-        show_plot: If True, displays the plot via ``plt.show()``.
-        show_colorbar: If True, displays a colorbar next to the plot (or a colorbar next to every submatrix if ).
-        scale_colors_per_submatrix: If True, considers the minimum and maximum of each Jacobian submatrix separately for color scaling. Otherwise, uses the global minimum and maximum of all Jacobians.
-        title: The title of the plot (optional).
-        colormap: The colormap to use for the plot.
-        log_scale: If True, uses a logarithmic scale for the matrix values shown in the image plot.
-
-    Returns:
-        The created Matplotlib figure.
-    """
-    wp.utils.warn(
-        "The function `plot_kernel_jacobians` is deprecated and will be removed in a future Warp version. Please call `jacobian_plot` instead.",
-        DeprecationWarning,
-        stacklevel=3,
-    )
-    return jacobian_plot(
-        jacobians,
-        kernel,
-        inputs,
-        outputs,
-        show_plot=show_plot,
-        show_colorbar=show_colorbar,
-        scale_colors_per_submatrix=scale_colors_per_submatrix,
-        title=title,
-        colormap=colormap,
-        log_scale=log_scale,
-    )
-
-
 def scalarize_array_1d(arr):
     # convert array to 1D array with scalar dtype
     if arr.dtype in wp.types.scalar_types:
@@ -724,19 +665,18 @@ def scalarize_array_2d(arr):
 
 
 def jacobian(
-    function: Union[wp.Kernel, Callable],
-    dim: Tuple[int] = None,
-    inputs: Sequence = None,
-    outputs: Sequence = None,
-    input_output_mask: List[Tuple[Union[str, int], Union[str, int]]] = None,
+    function: wp.Kernel | Callable,
+    dim: tuple[int] | None = None,
+    inputs: Sequence | None = None,
+    outputs: Sequence | None = None,
+    input_output_mask: list[tuple[str | int, str | int]] | None = None,
     device: wp.context.Devicelike = None,
     max_blocks=0,
     block_dim=256,
     max_outputs_per_var=-1,
     plot_jacobians=False,
-    metadata: FunctionMetadata = None,
-    kernel: wp.Kernel = None,
-) -> Dict[Tuple[int, int], wp.array]:
+    metadata: FunctionMetadata | None = None,
+) -> dict[tuple[int, int], wp.array]:
     """
     Computes the Jacobians of a function or Warp kernel for the provided selection of differentiable inputs to differentiable outputs.
 
@@ -764,20 +704,12 @@ def jacobian(
         max_outputs_per_var: Maximum number of output dimensions over which to evaluate the Jacobians for the input-output pairs. Evaluates all output dimensions if value <= 0.
         plot_jacobians: If True, visualizes the computed Jacobians in a plot (requires ``matplotlib``).
         metadata: The metadata of the kernel function, containing the input and output labels, strides, and dtypes. If None or empty, the metadata is inferred from the kernel or function.
-        kernel: Deprecated argument. Use the ``function`` argument instead.
 
     Returns:
         A dictionary of Jacobians, where the keys are tuples of input and output indices, and the values are the Jacobian matrices.
     """
     if input_output_mask is None:
         input_output_mask = []
-    if kernel is not None:
-        wp.utils.warn(
-            "The argument `kernel` to the function `wp.autograd.jacobian` is deprecated in favor of the `function` argument and will be removed in a future Warp version.",
-            DeprecationWarning,
-            stacklevel=3,
-        )
-        function = kernel
 
     if metadata is None:
         metadata = FunctionMetadata()
@@ -868,20 +800,19 @@ def jacobian(
 
 
 def jacobian_fd(
-    function: Union[wp.Kernel, Callable],
-    dim: Tuple[int] = None,
-    inputs: Sequence = None,
-    outputs: Sequence = None,
-    input_output_mask: List[Tuple[Union[str, int], Union[str, int]]] = None,
+    function: wp.Kernel | Callable,
+    dim: tuple[int] | None | None = None,
+    inputs: Sequence | None = None,
+    outputs: Sequence | None = None,
+    input_output_mask: list[tuple[str | int, str | int]] | None = None,
     device: wp.context.Devicelike = None,
     max_blocks=0,
     block_dim=256,
     max_inputs_per_var=-1,
     eps: float = 1e-4,
     plot_jacobians=False,
-    metadata: FunctionMetadata = None,
-    kernel: wp.Kernel = None,
-) -> Dict[Tuple[int, int], wp.array]:
+    metadata: FunctionMetadata | None = None,
+) -> dict[tuple[int, int], wp.array]:
     """
     Computes the finite-difference Jacobian of a function or Warp kernel for the provided selection of differentiable inputs to differentiable outputs.
     The method uses a central difference scheme to approximate the Jacobian.
@@ -911,20 +842,12 @@ def jacobian_fd(
         eps: The finite-difference step size.
         plot_jacobians: If True, visualizes the computed Jacobians in a plot (requires ``matplotlib``).
         metadata: The metadata of the kernel function, containing the input and output labels, strides, and dtypes. If None or empty, the metadata is inferred from the kernel or function.
-        kernel: Deprecated argument. Use the ``function`` argument instead.
 
     Returns:
         A dictionary of Jacobians, where the keys are tuples of input and output indices, and the values are the Jacobian matrices.
     """
     if input_output_mask is None:
         input_output_mask = []
-    if kernel is not None:
-        wp.utils.warn(
-            "The argument `kernel` to the function `wp.autograd.jacobian` is deprecated in favor of the `function` argument and will be removed in a future Warp version.",
-            DeprecationWarning,
-            stacklevel=3,
-        )
-        function = kernel
 
     if metadata is None:
         metadata = FunctionMetadata()
@@ -1003,8 +926,8 @@ def jacobian_fd(
         outputs_until_right = [conditional_clone(output) for output in outputs_copy[:output_i]]
         outputs_after_left = [conditional_clone(output) for output in outputs_copy[output_i + 1 :]]
         outputs_after_right = [conditional_clone(output) for output in outputs_copy[output_i + 1 :]]
-        left_outputs = outputs_until_left + [left] + outputs_after_left
-        right_outputs = outputs_until_right + [right] + outputs_after_right
+        left_outputs = [*outputs_until_left, left, *outputs_after_left]
+        right_outputs = [*outputs_until_right, right, *outputs_after_right]
 
         input_num = flat_input.shape[0]
         flat_input_copy = wp.clone(flat_input)