Basic padding functionality

Author: jessegrabowski

pytensor/tensor/basic.py

@@ -16,6 +16,7 @@
 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
@@ -50,6 +51,7 @@
     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,
@@ -61,6 +63,7 @@
     shape_tuple,
     specify_broadcastable,
 )
+from pytensor.tensor.subtensor import set_subtensor
 from pytensor.tensor.type import (
     TensorType,
     discrete_dtypes,
@@ -4342,6 +4345,125 @@ 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",