From 4a5992c7023d94c0b9fcbaaa7f1bd4eb1d03b126 Mon Sep 17 00:00:00 2001 From: Anton Volkov Date: Wed, 23 Jul 2025 18:04:52 +0200 Subject: [PATCH 1/3] Support pad_width as a dictionary --- dpnp/dpnp_utils/dpnp_utils_pad.py | 53 ++++++++++++++++++++++++++++--- 1 file changed, 49 insertions(+), 4 deletions(-) diff --git a/dpnp/dpnp_utils/dpnp_utils_pad.py b/dpnp/dpnp_utils/dpnp_utils_pad.py index 65f314bbf2fd..c3f3d7b34590 100644 --- a/dpnp/dpnp_utils/dpnp_utils_pad.py +++ b/dpnp/dpnp_utils/dpnp_utils_pad.py @@ -273,6 +273,47 @@ def _get_stats(padded, axis, width_pair, length_pair, stat_func): return left_stat, right_stat +def _pad_normalize_dict_width(pad_width, ndim): + """ + Normalize pad width passed as a dictionary. + + Parameters + ---------- + pad_width : dict + Padding specification. The keys must be integer axis indices, and + the values must be either: + - a single int (same padding before and after), + - a tuple of two ints (before, after). + ndim : int + Number of dimensions in the input array. + + Returns + ------- + seq : list + A (ndim, 2) list of padding widths for each axis. + + Raises + ------ + TypeError + If the padding format for any axis is invalid. + + """ + + seq = [(0, 0)] * ndim + for axis, width in pad_width.items(): + if isinstance(width, int): + seq[axis] = (width, width) + elif ( + isinstance(width, tuple) + and len(width) == 2 + and all(isinstance(w, int) for w in width) + ): + seq[axis] = width + else: + raise TypeError(f"Invalid pad width for axis {axis}: {width}") + return seq + + def _pad_simple(array, pad_width, fill_value=None): """ Copied from numpy/lib/_arraypad_impl.py @@ -616,21 +657,25 @@ def _view_roi(array, original_area_slice, axis): def dpnp_pad(array, pad_width, mode="constant", **kwargs): """Pad an array.""" + nd = array.ndim + if isinstance(pad_width, int): if pad_width < 0: raise ValueError("index can't contain negative values") - pad_width = ((pad_width, pad_width),) * array.ndim + pad_width = ((pad_width, pad_width),) * nd else: if dpnp.is_supported_array_type(pad_width): pad_width = dpnp.asnumpy(pad_width) else: + if isinstance(pad_width, dict): + pad_width = _pad_normalize_dict_width(pad_width, nd) pad_width = numpy.asarray(pad_width) if not pad_width.dtype.kind == "i": raise TypeError("`pad_width` must be of integral type.") - # Broadcast to shape (array.ndim, 2) - pad_width = _as_pairs(pad_width, array.ndim, as_index=True) + # Broadcast to shape (nd, 2) + pad_width = _as_pairs(pad_width, nd, as_index=True) if callable(mode): function = mode @@ -683,7 +728,7 @@ def dpnp_pad(array, pad_width, mode="constant", **kwargs): if ( dpnp.isscalar(values) and values == 0 - and (array.ndim == 1 or array.size < 3e7) + and (nd == 1 or array.size < 3e7) ): # faster path for 1d arrays or small n-dimensional arrays return _pad_simple(array, pad_width, 0)[0] From 474100b9d3be1b4e6d83dd23924530bff0e7c1a2 Mon Sep 17 00:00:00 2001 From: Anton Volkov Date: Wed, 23 Jul 2025 18:12:53 +0200 Subject: [PATCH 2/3] Update the docstring and add test --- dpnp/dpnp_iface_manipulation.py | 25 ++++++++++++++++++++++++- dpnp/tests/test_arraypad.py | 18 ++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/dpnp/dpnp_iface_manipulation.py b/dpnp/dpnp_iface_manipulation.py index 2806f621b5c6..464e99df2b64 100644 --- a/dpnp/dpnp_iface_manipulation.py +++ b/dpnp/dpnp_iface_manipulation.py @@ -2504,7 +2504,7 @@ def pad(array, pad_width, mode="constant", **kwargs): ---------- array : {dpnp.ndarray, usm_ndarray} The array of rank ``N`` 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. @@ -2512,6 +2512,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 dictionary, each key is an axis and its corresponding value is an + integer or a pair of integers describing the padding ``(before, after)`` + or ``pad`` width for that axis. mode : {str, function}, optional One of the following string values or a user supplied function. @@ -2694,6 +2697,26 @@ def pad(array, pad_width, mode="constant", **kwargs): [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]]) + """ dpnp.check_supported_arrays_type(array) diff --git a/dpnp/tests/test_arraypad.py b/dpnp/tests/test_arraypad.py index 795311105562..9a88dd8bab96 100644 --- a/dpnp/tests/test_arraypad.py +++ b/dpnp/tests/test_arraypad.py @@ -539,3 +539,21 @@ def test_as_pairs_exceptions(self): dpnp_as_pairs([[1, 2], [3, 4]], 3) with pytest.raises(ValueError, match="could not be broadcast"): dpnp_as_pairs(dpnp.ones((2, 3)), 3) + + @testing.with_requires("numpy>=2.4") + @pytest.mark.parametrize( + "sh, pad_width", + [ + ((3, 4, 5), {-2: (1, 3)}), + ((3, 4, 5), {0: (5, 2)}), + ((3, 4, 5), {0: (5, 2), -1: (3, 4)}), + ((3, 4, 5), {1: 5}), + ], + ) + def test_dict_pad_width(self, sh, pad_width): + a = numpy.zeros(sh) + ia = dpnp.array(a) + + result = dpnp.pad(ia, pad_width) + expected = numpy.pad(a, pad_width) + assert_equal(result, expected) From 0b1fd1c3474ca3441ac47bd10b5d9911f4193035 Mon Sep 17 00:00:00 2001 From: Anton Volkov Date: Wed, 23 Jul 2025 18:16:23 +0200 Subject: [PATCH 3/3] Add PR to the changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index ab8d1fd619e8..8e26886781ae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * Removed the use of class template argument deduction for alias template to conform to the C++17 standard [#2517](https://github.com/IntelPython/dpnp/pull/2517) * Changed th order of individual FFTs over `axes` for `dpnp.fft.irfftn` to be in forward order [#2524](https://github.com/IntelPython/dpnp/pull/2524) * Replaced the use of `numpy.testing.suppress_warnings` with appropriate calls from the warnings module [#2529](https://github.com/IntelPython/dpnp/pull/2529) +* Extended `dpnp.pad` to support `pad_width` keyword as a dictionary [#2535](https://github.com/IntelPython/dpnp/pull/2535) ### Deprecated