diff --git a/doc/release/upcoming_changes/29273.new_feature.rst b/doc/release/upcoming_changes/29273.new_feature.rst
new file mode 100644
index 000000000000..3e380ca0dbe6
--- /dev/null
+++ 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
index 507a0ab51b52..681b92fc8a72 100644
--- 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,7 +552,7 @@ 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.
@@ -558,6 +560,9 @@ def pad(array, pad_width, mode='constant', **kwargs):
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
index 6158f299efe2..cfabdecf669e 100644
--- 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
index 6efbe348ca81..14383e743e47 100644
--- 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
index c5a443d93fe3..3d53d913a770 100644
--- a/numpy/typing/tests/data/reveal/arraypad.pyi
+++ b/numpy/typing/tests/data/reveal/arraypad.pyi
@@ -20,3 +20,8 @@ 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)}), npt.NDArray[np.int64])
+assert_type(np.pad(AR_i8, {-2: 4}), npt.NDArray[np.int64])
+pad_width: dict[int, int | tuple[int, int]] = {-1: (2, 3), -2: 4}
+assert_type(np.pad(AR_i8, pad_width), npt.NDArray[np.int64])
pFad - Phonifier reborn
Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.
Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.
Alternative Proxies:
Alternative Proxy
pFad Proxy
pFad v3 Proxy
pFad v4 Proxy