ENH: Extend numpy.pad to handle pad_width dictionary argument.

carlosgmartin · carlosgmartin · commit cb021ebe677a · 2025-06-26T08:12:32.000-04:00
diff --git a/doc/release/upcoming_changes/29273.new_feature.rst b/doc/release/upcoming_changes/29273.new_feature.rst
@@ -0,0 +1 @@
+Extend ``numpy.pad`` to accept a dictionary for the ``pad_width`` argument.
diff --git a/numpy/lib/_arraypad_impl.py b/numpy/lib/_arraypad_impl.py
@@ -3,6 +3,8 @@
 of an n-dimensional array.
 
 """
+import typing
+
 import numpy as np
 from numpy._core.overrides import array_function_dispatch
 from numpy.lib._index_tricks_impl import ndindex
@@ -550,14 +552,17 @@ def pad(array, pad_width, mode='constant', **kwargs):
     ----------
     array : array_like of rank N
         The array to pad.
-    pad_width : {sequence, array_like, int}
+    pad_width : {sequence, array_like, int, dict}
         Number of values padded to the edges of each axis.
         ``((before_1, after_1), ... (before_N, after_N))`` unique pad widths
         for each axis.
         ``(before, after)`` or ``((before, after),)`` yields same before
         and after pad for each axis.
         ``(pad,)`` or ``int`` is a shortcut for before = after = pad width
         for all axes.
+        If a ``dict``, each key is an axis and its corresponding value is an ``int`` or
+        ``int`` pair describing the padding ``(before, after)`` or ``pad`` width for
+        that axis.
     mode : str or function, optional
         One of the following string values or a user supplied function.
 
@@ -745,8 +750,39 @@ def pad(array, pad_width, mode='constant', **kwargs):
            [100, 100,   3,   4,   5, 100, 100],
            [100, 100, 100, 100, 100, 100, 100],
            [100, 100, 100, 100, 100, 100, 100]])
+
+    >>> a = np.arange(1, 7).reshape(2, 3)
+    >>> np.pad(a, {1: (1, 2)})
+    array([[0, 1, 2, 3, 0, 0],
+           [0, 4, 5, 6, 0, 0]])
+    >>> np.pad(a, {-1: 2})
+    array([[0, 0, 1, 2, 3, 0, 0],
+           [0, 0, 4, 5, 6, 0, 0]])
+    >>> np.pad(a, {0: (3, 0)})
+    array([[0, 0, 0],
+           [0, 0, 0],
+           [0, 0, 0],
+           [1, 2, 3],
+           [4, 5, 6]])
+    >>> np.pad(a, {0: (3, 0), 1: 2})
+    array([[0, 0, 0, 0, 0, 0, 0],
+           [0, 0, 0, 0, 0, 0, 0],
+           [0, 0, 0, 0, 0, 0, 0],
+           [0, 0, 1, 2, 3, 0, 0],
+           [0, 0, 4, 5, 6, 0, 0]])
     """
     array = np.asarray(array)
+    if isinstance(pad_width, dict):
+        seq = [(0, 0)] * array.ndim
+        for axis, width in pad_width.items():
+            match width:
+                case int(both):
+                    seq[axis] = both, both
+                case tuple((int(before), int(after))):
+                    seq[axis] = before, after
+                case _ as invalid:
+                    typing.assert_never(invalid)
+        pad_width = seq
     pad_width = np.asarray(pad_width)
 
     if not pad_width.dtype.kind == 'i':
diff --git a/numpy/lib/_arraypad_impl.pyi b/numpy/lib/_arraypad_impl.pyi
@@ -43,11 +43,17 @@ _ModeKind: TypeAlias = L[
 # TODO: In practice each keyword argument is exclusive to one or more
 # specific modes. Consider adding more overloads to express this in the future.
 
+_PadWidth: TypeAlias = (
+    _ArrayLikeInt
+    | dict[int, int]
+    | dict[int, tuple[int, int]]
+    | dict[int, int | tuple[int, int]]
+)
 # Expand `**kwargs` into explicit keyword-only arguments
 @overload
 def pad(
     array: _ArrayLike[_ScalarT],
-    pad_width: _ArrayLikeInt,
+    pad_width: _PadWidth,
     mode: _ModeKind = ...,
     *,
     stat_length: _ArrayLikeInt | None = ...,
@@ -58,7 +64,7 @@ def pad(
 @overload
 def pad(
     array: ArrayLike,
-    pad_width: _ArrayLikeInt,
+    pad_width: _PadWidth,
     mode: _ModeKind = ...,
     *,
     stat_length: _ArrayLikeInt | None = ...,
@@ -69,14 +75,14 @@ def pad(
 @overload
 def pad(
     array: _ArrayLike[_ScalarT],
-    pad_width: _ArrayLikeInt,
+    pad_width: _PadWidth,
     mode: _ModeFunc,
     **kwargs: Any,
 ) -> NDArray[_ScalarT]: ...
 @overload
 def pad(
     array: ArrayLike,
-    pad_width: _ArrayLikeInt,
+    pad_width: _PadWidth,
     mode: _ModeFunc,
     **kwargs: Any,
 ) -> NDArray[Any]: ...
diff --git a/numpy/lib/tests/test_arraypad.py b/numpy/lib/tests/test_arraypad.py
@@ -1413,3 +1413,15 @@ def test_dtype_persistence(dtype, mode):
     arr = np.zeros((3, 2, 1), dtype=dtype)
     result = np.pad(arr, 1, mode=mode)
     assert result.dtype == dtype
+
+
+@pytest.mark.parametrize("input_shape, pad_width, expected_shape", [
+    ((3, 4, 5), {-2: (1, 3)}, (3, 4 + 1 + 3, 5)),
+    ((3, 4, 5), {0: (5, 2)}, (3 + 5 + 2, 4, 5)),
+    ((3, 4, 5), {0: (5, 2), -1: (3, 4)}, (3 + 5 + 2, 4, 5 + 3 + 4)),
+    ((3, 4, 5), {1: 5}, (3, 4 + 2 * 5, 5)),
+])
+def test_pad_dict_pad_width(input_shape, pad_width, expected_shape):
+    a = np.zeros(input_shape)
+    result = np.pad(a, pad_width)
+    assert result.shape == expected_shape
diff --git a/numpy/typing/tests/data/reveal/arraypad.pyi b/numpy/typing/tests/data/reveal/arraypad.pyi
@@ -20,3 +20,7 @@ assert_type(np.pad(AR_LIKE, (2, 3), "constant"), npt.NDArray[Any])
 
 assert_type(np.pad(AR_f8, (2, 3), mode_func), npt.NDArray[np.float64])
 assert_type(np.pad(AR_f8, (2, 3), mode_func, a=1, b=2), npt.NDArray[np.float64])
+
+assert_type(np.pad(AR_i8, {-1: (2, 3)}, "constant"), npt.NDArray[np.int64])
+assert_type(np.pad(AR_i8, {-2: 4}, "constant"), npt.NDArray[np.int64])
+assert_type(np.pad(AR_f8, {-1: (2, 3), -2: 4}, "constant"), npt.NDArray[np.float64])

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+Extend ``numpy.pad`` to accept a dictionary for the ``pad_width`` argument.