move pad functions to pad.py

Author: jessegrabowski

pytensor/tensor/__init__.py

@@ -130,6 +130,7 @@ def _get_vector_length_Constant(op: Op | Variable, var: Constant) -> int:
 from pytensor.tensor.extra_ops import *
 from pytensor.tensor.io import *
 from pytensor.tensor.math import *
+from pytensor.tensor.pad import pad
 from pytensor.tensor.shape import (
     reshape,
     shape,

pytensor/tensor/basic.py

@@ -16,7 +16,6 @@
 import numpy as np
 from numpy.core.multiarray import normalize_axis_index
 from numpy.core.numeric import normalize_axis_tuple
-from numpy.lib.arraypad import _get_edges, _slice_at_axis
 
 import pytensor
 import pytensor.scalar.sharedvar
@@ -51,7 +50,6 @@
     scalar_elemwise,
 )
 from pytensor.tensor.exceptions import NotScalarConstantError
-from pytensor.tensor.extra_ops import broadcast_to, linspace
 from pytensor.tensor.shape import (
     Shape,
     Shape_i,
@@ -63,7 +61,6 @@
     shape_tuple,
     specify_broadcastable,
 )
-from pytensor.tensor.subtensor import set_subtensor
 from pytensor.tensor.type import (
     TensorType,
     discrete_dtypes,
@@ -4345,125 +4342,6 @@ def ix_(*args):
     return tuple(out)
 
 
-def _symbolic_pad(
-    x: TensorVariable, pad_width: TensorVariable
-) -> tuple[TensorVariable, tuple[slice, ...], TensorVariable]:
-    pad_width = broadcast_to(pad_width, (x.ndim, 2))
-    new_shape = as_tensor(
-        [pad_width[i][0] + size + pad_width[i][1] for i, size in enumerate(x.shape)]
-    )
-    original_area_slice = tuple(
-        slice(pad_width[i][0], pad_width[i][0] + size) for i, size in enumerate(x.shape)
-    )
-    padded: TensorVariable = set_subtensor(zeros(new_shape)[original_area_slice], x)
-    return padded, original_area_slice, pad_width
-
-
-def _get_padding_slices(
-    dim_shape: TensorVariable, width_pair: tuple[TensorVariable], axis: int
-):
-    left_slice = _slice_at_axis(slice(None, width_pair[0]), axis)
-    right_slice = _slice_at_axis(slice(dim_shape - width_pair[1], None), axis)
-
-    return left_slice, right_slice
-
-
-def _constant_pad(
-    x: TensorVariable, pad_width: TensorVariable, constant_values: TensorVariable
-):
-    padded, area_slice, pad_width = _symbolic_pad(x, pad_width)
-    values = broadcast_to(constant_values, (padded.ndim, 2))
-
-    for axis in range(padded.ndim):
-        width_pair = pad_width[axis]
-        value_pair = values[axis]
-        dim_shape = padded.shape[axis]
-
-        left_slice, right_slice = _get_padding_slices(dim_shape, width_pair, axis)
-        padded = set_subtensor(padded[left_slice], value_pair[0])
-        padded = set_subtensor(padded[right_slice], value_pair[1])
-
-    return padded
-
-
-def _edge_pad(x: TensorVariable, pad_width: TensorVariable):
-    padded, area_slice, pad_width = _symbolic_pad(x, pad_width)
-    for axis in range(padded.ndim):
-        width_pair = pad_width[axis]
-        dim_shape = padded.shape[axis]
-
-        left_edge, right_edge = _get_edges(padded, axis, width_pair)
-        left_slice, right_slice = _get_padding_slices(dim_shape, width_pair, axis)
-
-        padded = set_subtensor(padded[left_slice], left_edge)
-        padded = set_subtensor(padded[right_slice], right_edge)
-
-    return padded
-
-
-def _linear_ramp_pad(
-    x: TensorVariable, pad_width: TensorVariable, end_values: TensorVariable = 0
-):
-    padded, area_slice, pad_width = _symbolic_pad(x, pad_width)
-    end_values = broadcast_to(end_values, (padded.ndim, 2))
-    for axis in range(padded.ndim):
-        width_pair = pad_width[axis]
-        end_value_pair = end_values[axis]
-        edge_pair = _get_edges(padded, axis, width_pair)
-        dim_shape = padded.shape[axis]
-        left_slice, right_slice = _get_padding_slices(dim_shape, width_pair, axis)
-
-        # pt.linspace doesn't have the endpoint kwarg, so need to take one extra step then slice it away
-        left_ramp = linspace(
-            start=end_value_pair[0],
-            end=specify_broadcastable(edge_pair[0], axis).squeeze(axis),
-            steps=width_pair[0] + 1,
-        )[:-1]
-        right_ramp = linspace(
-            start=end_value_pair[1],
-            end=specify_broadcastable(edge_pair[1], axis).squeeze(axis),
-            steps=width_pair[1] + 1,
-        )[:-1]
-        right_ramp = right_ramp[_slice_at_axis(slice(None, None, -1), axis)]
-
-        # FIXME: This swapaxes is needed because the shapes of the linspaces don't "rotate" with
-        #  the different dimensions. But this makes the non-active dimensions backwards in the padding.
-        padded = set_subtensor(padded[left_slice], swapaxes(left_ramp, 0, axis))
-        padded = set_subtensor(padded[right_slice], swapaxes(right_ramp, 0, axis))
-
-    return padded
-
-
-def pad(x, pad_width, mode="constant", **kwargs):
-    allowed_kwargs = {
-        "empty": [],
-        "edge": [],
-        "wrap": [],
-        "constant": ["constant_values"],
-        "linear_ramp": ["end_values"],
-        "maximum": ["stat_length"],
-        "mean": ["stat_length"],
-        "median": ["stat_length"],
-        "minimum": ["stat_length"],
-        "reflect": ["reflect_type"],
-        "symmetric": ["reflect_type"],
-    }
-
-    if any(value not in allowed_kwargs[mode] for value in kwargs.values()):
-        raise ValueError(
-            f"Invalid keyword arguments for mode '{mode}': {kwargs.keys()}"
-        )
-
-    if mode == "constant":
-        constant_values = kwargs.pop("constant_values", 0)
-        return _constant_pad(x, pad_width, constant_values)
-    elif mode == "edge":
-        return _edge_pad(x, pad_width)
-    elif mode == "linear_ramp":
-        end_values = kwargs.pop("end_values", 0)
-        return _linear_ramp_pad(x, pad_width, end_values)
-
-
 __all__ = [
     "take_along_axis",
     "expand_dims",

pytensor/tensor/pad.py

@@ -0,0 +1,133 @@
+from numpy.lib.arraypad import _get_edges, _slice_at_axis  # noqa
+
+from pytensor.tensor.basic import (
+    TensorVariable,
+    as_tensor,
+    swapaxes,
+    zeros,
+)
+from pytensor.tensor.extra_ops import linspace, broadcast_to
+from pytensor.tensor.shape import specify_broadcastable
+from pytensor.tensor.subtensor import set_subtensor
+
+
+def _symbolic_pad(
+    x: TensorVariable, pad_width: TensorVariable
+) -> tuple[TensorVariable, tuple[slice, ...], TensorVariable]:
+    pad_width = broadcast_to(pad_width, (x.ndim, 2))
+    new_shape = as_tensor(
+        [pad_width[i][0] + size + pad_width[i][1] for i, size in enumerate(x.shape)]
+    )
+    original_area_slice = tuple(
+        slice(pad_width[i][0], pad_width[i][0] + size) for i, size in enumerate(x.shape)
+    )
+    padded: TensorVariable = set_subtensor(zeros(new_shape)[original_area_slice], x)
+    return padded, original_area_slice, pad_width
+
+
+def _get_padding_slices(
+    dim_shape: TensorVariable, width_pair: tuple[TensorVariable, TensorVariable], axis: int
+):
+    left_slice = _slice_at_axis(slice(None, width_pair[0]), axis)
+    right_slice = _slice_at_axis(slice(dim_shape - width_pair[1], None), axis)
+
+    return left_slice, right_slice
+
+
+def _constant_pad(
+    x: TensorVariable, pad_width: TensorVariable, constant_values: TensorVariable
+):
+    padded, area_slice, pad_width = _symbolic_pad(x, pad_width)
+    values = broadcast_to(constant_values, (padded.ndim, 2))
+
+    for axis in range(padded.ndim):
+        width_pair = pad_width[axis]
+        value_pair = values[axis]
+        dim_shape = padded.shape[axis]
+
+        left_slice, right_slice = _get_padding_slices(dim_shape, width_pair, axis)
+        padded = set_subtensor(padded[left_slice], value_pair[0])
+        padded = set_subtensor(padded[right_slice], value_pair[1])
+
+    return padded
+
+
+def _edge_pad(x: TensorVariable, pad_width: TensorVariable):
+    padded, area_slice, pad_width = _symbolic_pad(x, pad_width)
+    for axis in range(padded.ndim):
+        width_pair = pad_width[axis]
+        dim_shape = padded.shape[axis]
+
+        left_edge, right_edge = _get_edges(padded, axis, width_pair)
+        left_slice, right_slice = _get_padding_slices(dim_shape, width_pair, axis)
+
+        padded = set_subtensor(padded[left_slice], left_edge)
+        padded = set_subtensor(padded[right_slice], right_edge)
+
+    return padded
+
+
+def _linear_ramp_pad(
+    x: TensorVariable, pad_width: TensorVariable, end_values: TensorVariable | int = 0
+):
+    padded, area_slice, pad_width = _symbolic_pad(x, pad_width)
+    end_values = broadcast_to(end_values, (padded.ndim, 2))
+    for axis in range(padded.ndim):
+        width_pair = pad_width[axis]
+        end_value_pair = end_values[axis]
+        edge_pair = _get_edges(padded, axis, width_pair)
+        dim_shape = padded.shape[axis]
+        left_slice, right_slice = _get_padding_slices(dim_shape, width_pair, axis)
+
+        # pt.linspace doesn't have the endpoint kwarg, so need to take one extra step then slice it away
+        left_ramp = linspace(
+            start=end_value_pair[0],
+            end=specify_broadcastable(edge_pair[0], axis).squeeze(axis),
+            steps=width_pair[0] + 1,
+        )[:-1]
+        right_ramp = linspace(
+            start=end_value_pair[1],
+            end=specify_broadcastable(edge_pair[1], axis).squeeze(axis),
+            steps=width_pair[1] + 1,
+        )[:-1]
+        right_ramp = right_ramp[_slice_at_axis(slice(None, None, -1), axis)]
+
+        # FIXME: This swapaxes is needed because the shapes of the linspaces don't "rotate" with
+        #  the different dimensions. But this makes the non-active dimensions backwards in the padding.
+        padded = set_subtensor(padded[left_slice], swapaxes(left_ramp, 0, axis))
+        padded = set_subtensor(padded[right_slice], swapaxes(right_ramp, 0, axis))
+
+    return padded
+
+
+def pad(x, pad_width, mode="constant", **kwargs):
+    allowed_kwargs = {
+        "empty": [],
+        "edge": [],
+        "wrap": [],
+        "constant": ["constant_values"],
+        "linear_ramp": ["end_values"],
+        "maximum": ["stat_length"],
+        "mean": ["stat_length"],
+        "median": ["stat_length"],
+        "minimum": ["stat_length"],
+        "reflect": ["reflect_type"],
+        "symmetric": ["reflect_type"],
+    }
+
+    if any(value not in allowed_kwargs[mode] for value in kwargs.values()):
+        raise ValueError(
+            f"Invalid keyword arguments for mode '{mode}': {kwargs.keys()}"
+        )
+
+    if mode == "constant":
+        constant_values = kwargs.pop("constant_values", 0)
+        return _constant_pad(x, pad_width, constant_values)
+    elif mode == "edge":
+        return _edge_pad(x, pad_width)
+    elif mode == "linear_ramp":
+        end_values = kwargs.pop("end_values", 0)
+        return _linear_ramp_pad(x, pad_width, end_values)
+
+
+__all__ = ["pad"]