Implement pad

pytensor/tensor/pad.py

@@ -21,7 +21,7 @@
 from pytensor.tensor.math import max as pt_max
 from pytensor.tensor.math import min as pt_min
 from pytensor.tensor.shape import specify_broadcastable
-from pytensor.tensor.subtensor import flip, set_subtensor
+from pytensor.tensor.subtensor import flip, set_subtensor, slice_at_axis
 
 
 PadMode = Literal[
@@ -52,37 +52,6 @@
 }
 
 
-def _slice_at_axis(sl: slice, axis: int) -> tuple[slice, ...]:
-    """
-    Construct tuple of slices to slice an array in the given dimension.
-
-    Copied from numpy.lib.arraypad._slice_at_axis
-    https://github.com/numpy/numpy/blob/300096d384046eee479b0c7a70f79e308da52bff/numpy/lib/_arraypad_impl.py#L33
-
-    Parameters
-    ----------
-    sl : slice
-        The slice for the given dimension.
-    axis : int
-        The axis to which `sl` is applied. All other dimensions are left
-        "unsliced".
-
-    Returns
-    -------
-    sl : tuple of slices
-        A tuple with slices matching `shape` in length.
-
-    Examples
-    --------
-
-    .. code-block:: python
-
-         _slice_at_axis(slice(None, 3, -1), 1)
-        (slice(None, None, None), slice(None, 3, -1), (...,))
-    """
-    return (slice(None),) * axis + (sl,) + (...,)  # type: ignore
-
-
 def _get_edges(
     padded: TensorVariable, axis: int, width_pair: tuple[TensorVariable, TensorVariable]
 ) -> tuple[TensorVariable, TensorVariable]:
@@ -110,11 +79,11 @@
         `axis` which will have a length of 1.
     """
     left_index = width_pair[0]
-    left_slice = _slice_at_axis(slice(left_index, left_index + 1), axis)
+    left_slice = slice_at_axis(slice(left_index, left_index + 1), axis)
     left_edge = padded[left_slice]
 
     right_index = padded.shape[axis] - width_pair[1]
-    right_slice = _slice_at_axis(slice(right_index - 1, right_index), axis)
+    right_slice = slice_at_axis(slice(right_index - 1, right_index), axis)
     right_edge = padded[right_slice]
 
     return left_edge, right_edge
@@ -139,8 +108,8 @@
     width_pair: tuple[TensorVariable, TensorVariable],
     axis: int,
 ) -> tuple[tuple[slice, ...], tuple[slice, ...]]:
-    left_slice = _slice_at_axis(slice(None, width_pair[0]), axis)
-    right_slice = _slice_at_axis(slice(dim_shape - width_pair[1], None), axis)
+    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
 
@@ -224,20 +193,20 @@
     left_length = (
         minimum(left_length, max_length) if left_length is not None else max_length
     )
-    left_slice = _slice_at_axis(slice(left_index, left_index + left_length), axis)
+    left_slice = slice_at_axis(slice(left_index, left_index + left_length), axis)
     left_chunk = padded[left_slice]
     left_stat = stat_func(left_chunk, axis=axis, keepdims=True)
     if left_length is None and right_length is None:
         # We could also return early in the more general case of left_length == right_length, but we don't necessarily
         # know these shapes.
         # TODO: Add rewrite to simplify in this case
         return left_stat, left_stat
 
     # Calculate statistic for the right side
     right_length = (
         minimum(right_length, max_length) if right_length is not None else max_length
     )
-    right_slice = _slice_at_axis(slice(right_index - right_length, right_index), axis)
+    right_slice = slice_at_axis(slice(right_index - right_length, right_index), axis)
     right_chunk = padded[right_slice]
     right_stat = stat_func(right_chunk, axis=axis, keepdims=True)
 
@@ -298,7 +267,7 @@
         )
 
         # Reverse the direction of the ramp for the "right" side
-        right_ramp = right_ramp[_slice_at_axis(slice(None, None, -1), axis)]  # type: ignore
+        right_ramp = right_ramp[slice_at_axis(slice(None, None, -1), axis)]  # type: ignore
 
         padded = set_subtensor(padded[left_slice], left_ramp)
         padded = set_subtensor(padded[right_slice], right_ramp)
@@ -336,7 +305,7 @@
         x = parts.reshape(new_shape)
 
         # Trim the excess on the active dimension
-        trim_slice = _slice_at_axis(slice(left_trim, -right_trim), axis)
+        trim_slice = slice_at_axis(slice(left_trim, -right_trim), axis)
         x = x[trim_slice]
 
     return x
@@ -393,7 +362,7 @@
         )
         right_trim = x.shape[axis] - right_trim
 
-        trim_slice = _slice_at_axis(slice(left_trim, right_trim), axis)
+        trim_slice = slice_at_axis(slice(left_trim, right_trim), axis)
         x = x[trim_slice]
 
     return x
@@ -407,7 +376,7 @@
     for axis in range(x.ndim):
         trimmed_size = x.shape[axis] - 1
 
-        trim_slice = _slice_at_axis(slice(None, -1), axis)
+        trim_slice = slice_at_axis(slice(None, -1), axis)
         x_trimmed = x[trim_slice]
         x_flipped = flip(x, axis=axis)[trim_slice]
 
@@ -429,8 +398,8 @@
             inner_func=partial(_reflect_inner, padding_left=False),
         )
 
-        left_trim = _slice_at_axis(slice(trimmed_size - remainders[0] - 1, -1), axis)
-        right_trim = _slice_at_axis(
+        left_trim = slice_at_axis(slice(trimmed_size - remainders[0] - 1, -1), axis)
+        right_trim = slice_at_axis(
             slice(1, right_padding.shape[axis] - trimmed_size + remainders[1] + 1), axis
         )
 
@@ -457,7 +426,7 @@
 
 def pad(x: TensorLike, pad_width: TensorLike, mode: PadMode = "constant", **kwargs):
     if any(value not in allowed_kwargs[mode] for value in kwargs.keys()):
         raise ValueError(
             f"Invalid keyword arguments for mode '{mode}': {kwargs.keys()}"
         )
     x = as_tensor(x, name="x")
@@ -479,7 +448,7 @@
         if mode == "median":
             # TODO: Revisit this after we implement a quantile function.
             #  See https://github.com/pymc-devs/pytensor/issues/53
             raise NotImplementedError("Median padding not implemented")
         stat_func = cast(Callable, stat_funcs[mode])
         stat_length = kwargs.get("stat_length")
         if stat_length is not None:

pytensor/tensor/subtensor.py

@@ -3021,8 +3021,57 @@ def _get_vector_length_Subtensor(op, var):
         raise ValueError(f"Length of {var} cannot be determined")
 
 
+def slice_at_axis(sl: slice, axis: int) -> tuple[slice, ...]:
+    """
+    Construct tuple of slices to slice an array in the given dimension.
+
+    Copied from numpy.lib.arraypad._slice_at_axis
+    https://github.com/numpy/numpy/blob/300096d384046eee479b0c7a70f79e308da52bff/numpy/lib/_arraypad_impl.py#L33
+
+    Parameters
+    ----------
+    sl : slice
+        The slice for the given dimension.
+    axis : int
+        The axis to which `sl` is applied. All other dimensions are left
+        "unsliced".
+
+    Returns
+    -------
+    sl : tuple of slices
+        A tuple with slices matching `shape` in length.
+
+    Examples
+    --------
+
+    .. code-block:: python
+
+        import pytensor.tensor as pt
+
+        s = pt.slice_at_axis(slice(None, 1), 1)
+        s
+        # Output: (slice(None, None, None), slice(None, 3, -1), (...,))
+
+        x = pt.tensor('x', shape=(None, None, None))
+        x_sliced = x[s]
+
+        f = pytensor.function([x], x_sliced)
+        x = np.arange(27).reshape(3, 3, 3)
+        f(x)
+        # Output: array([[[ 0.,  1.,  2.]],
+        #                [[ 9., 10., 11.]],
+        #                [[18., 19., 20.]]])
+    """
+    if axis >= 0:
+        return (slice(None),) * axis + (sl,) + (...,)  # type: ignore
+    else:
+        # If axis = -1 we want zero right padding (and so on), so subtract one
+        axis = abs(axis) - 1
+        return (...,) + (sl,) + (slice(None),) * axis  # type: ignore
+
+
 def flip(
-    arr: TensorVariable, axis: int | tuple[int] | TensorVariable = None
+    arr: TensorVariable, axis: int | tuple[int] | TensorVariable | None = None
 ) -> TensorVariable:
     """
     Reverse the order of elements in an tensor along the given axis.
@@ -3066,12 +3115,14 @@ def flip(
             slice(None, None, -1) if i in axis else slice(None, None, None)
             for i in range(arr.ndim)
         ]
+
     return arr[index]
 
 
 __all__ = [
     "take",
     "flip",
+    "slice_at_axis",
     "inc_subtensor",
     "set_subtensor",
 ]

tests/tensor/test_subtensor.py

@@ -43,6 +43,7 @@
     index_vars_to_types,
     indexed_result_shape,
     set_subtensor,
+    slice_at_axis,
     take,
 )
 from pytensor.tensor.type import (
@@ -2903,6 +2904,16 @@ def test_vectorize_adv_subtensor(
     )
 
 
+def test_slice_at_axis():
+    x = ptb.tensor("x", shape=(3, 4, 5))
+    x_sliced = x[slice_at_axis(slice(None, 1), axis=0)]
+    assert x_sliced.type.shape == (1, 4, 5)
+
+    # Negative axis
+    x_sliced = x[slice_at_axis(slice(None, 1), axis=-2)]
+    assert x_sliced.type.shape == (3, 1, 5)
+
+
 @pytest.mark.parametrize(
     "size", [(3,), (3, 3), (3, 5, 5)], ids=["1d", "2d square", "3d square"]
 )