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

carlosgmartin · carlosgmartin · commit 5ee3f62bbde3 · 2025-06-25T21:31:09.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,16 @@ 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 width for that axis.
     mode : str or function, optional
         One of the following string values or a user supplied function.
 
@@ -745,8 +749,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(2 * 3).reshape(2, 3)
+    >>> np.pad(a, {1: (1, 2)})
+    array([[0, 0, 1, 2, 0, 0],
+           [0, 3, 4, 5, 0, 0]])
+    >>> np.pad(a, {-1: 2})
+    array([[0, 0, 0, 1, 2, 0, 0],
+           [0, 0, 3, 4, 5, 0, 0]])
+    >>> np.pad(a, {0: (3, 0)})
+    array([[0, 0, 0],
+           [0, 0, 0],
+           [0, 0, 0],
+           [0, 1, 2],
+           [3, 4, 5]])
+    >>> 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, 0, 1, 2, 0, 0],
+           [0, 0, 3, 4, 5, 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
@@ -47,7 +47,7 @@ _ModeKind: TypeAlias = L[
 @overload
 def pad(
     array: _ArrayLike[_ScalarT],
-    pad_width: _ArrayLikeInt,
+    pad_width: _ArrayLikeInt | dict,
     mode: _ModeKind = ...,
     *,
     stat_length: _ArrayLikeInt | None = ...,
@@ -58,7 +58,7 @@ def pad(
 @overload
 def pad(
     array: ArrayLike,
-    pad_width: _ArrayLikeInt,
+    pad_width: _ArrayLikeInt | dict,
     mode: _ModeKind = ...,
     *,
     stat_length: _ArrayLikeInt | None = ...,
@@ -69,14 +69,14 @@ def pad(
 @overload
 def pad(
     array: _ArrayLike[_ScalarT],
-    pad_width: _ArrayLikeInt,
+    pad_width: _ArrayLikeInt | dict,
     mode: _ModeFunc,
     **kwargs: Any,
 ) -> NDArray[_ScalarT]: ...
 @overload
 def pad(
     array: ArrayLike,
-    pad_width: _ArrayLikeInt,
+    pad_width: _ArrayLikeInt | dict,
     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

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.