Merge pull request #14209 from lagru/fix-linear-ramp

BUG: Fix uint-overflow if padding with linear_ramp and negative gain

1 files changed, 16 insertions, 68 deletions

diff --git a/numpy/lib/arraypad.py b/numpy/lib/arraypad.py
index 62330e692..33e64708d 100644
--- a/numpy/lib/arraypad.py
+++ b/numpy/lib/arraypad.py
@@ -17,66 +17,6 @@ __all__ = ['pad']
 # Private utility functions.
 
 
-def _linear_ramp(ndim, axis, start, stop, size, reverse=False):
-    """
-    Create a linear ramp of `size` in `axis` with `ndim`.
-
-    This algorithm behaves like a vectorized version of `numpy.linspace`.
-    The resulting linear ramp is broadcastable to any array that matches the
-    ramp in `shape[axis]` and `ndim`.
-
-    Parameters
-    ----------
-    ndim : int
-        Number of dimensions of the resulting array. All dimensions except
-        the one specified by `axis` will have the size 1.
-    axis : int
-        The dimension that contains the linear ramp of `size`.
-    start : int or ndarray
-        The starting value(s) of the linear ramp. If given as an array, its
-        size must match `size`.
-    stop : int or ndarray
-        The stop value(s) (not included!) of the linear ramp. If given as an
-        array, its size must match `size`.
-    size : int
-        The number of elements in the linear ramp. If this argument is 0 the
-        dimensions of `ramp` will all be of length 1 except for the one given
-        by `axis` which will be 0.
-    reverse : bool
-        If False, increment in a positive fashion, otherwise decrement.
-
-    Returns
-    -------
-    ramp : ndarray
-        Output array of dtype np.float64 that in- or decrements along the given
-        `axis`.
-
-    Examples
-    --------
-    >>> _linear_ramp(ndim=2, axis=0, start=np.arange(3), stop=10, size=2)
-    array([[0. , 1. , 2. ],
-           [5. , 5.5, 6. ]])
-    >>> _linear_ramp(ndim=3, axis=0, start=2, stop=0, size=0)
-    array([], shape=(0, 1, 1), dtype=float64)
-    """
-    # Create initial ramp
-    ramp = np.arange(size, dtype=np.float64)
-    if reverse:
-        ramp = ramp[::-1]
-
-    # Make sure, that ramp is broadcastable
-    init_shape = (1,) * axis + (size,) + (1,) * (ndim - axis - 1)
-    ramp = ramp.reshape(init_shape)
-
-    if size != 0:
-        # And scale to given start and stop values
-        gain = (stop - start) / float(size)
-        ramp = ramp * gain
-        ramp += start
-
-    return ramp
-
-
 def _round_if_needed(arr, dtype):
     """
     Rounds arr inplace if destination dtype is integer.
@@ -269,17 +209,25 @@ def _get_linear_ramps(padded, axis, width_pair, end_value_pair):
     """
     edge_pair = _get_edges(padded, axis, width_pair)
 
-    left_ramp = _linear_ramp(
-        padded.ndim, axis, start=end_value_pair[0], stop=edge_pair[0],
-        size=width_pair[0], reverse=False
+    left_ramp = np.linspace(
+        start=end_value_pair[0],
+        stop=edge_pair[0].squeeze(axis),  # Dimensions is replaced by linspace
+        num=width_pair[0],
+        endpoint=False,
+        dtype=padded.dtype,
+        axis=axis,
     )
-    _round_if_needed(left_ramp, padded.dtype)
 
-    right_ramp = _linear_ramp(
-        padded.ndim, axis, start=end_value_pair[1], stop=edge_pair[1],
-        size=width_pair[1], reverse=True
+    right_ramp = np.linspace(
+        start=end_value_pair[1],
+        stop=edge_pair[1].squeeze(axis),  # Dimension is replaced by linspace
+        num=width_pair[1],
+        endpoint=False,
+        dtype=padded.dtype,
+        axis=axis,
     )
-    _round_if_needed(right_ramp, padded.dtype)
+    # Reverse linear space in appropriate dimension
+    right_ramp = right_ramp[_slice_at_axis(slice(None, None, -1), axis)]
 
     return left_ramp, right_ramp