Repository: numba/numba

Issue:
FR: Support `numpy.lib.stride_tricks.sliding_window_view`
Hi.
Please add support for numpy's `sliding_window_view` ([docs](https://numpy.org/doc/stable/reference/generated/numpy.lib.stride_tricks.sliding_window_view.html#numpy.lib.stride_tricks.sliding_window_view)). It is essentially a safe and useful wrapper to `.as_strided()` that you already support so that shouldn't be much of an issue. 

Thanks!


Discussion:
Many thanks for the request! I'll label this issue as such.

Resolving patch:
diff --git a/numba/np/arrayobj.py b/numba/np/arrayobj.py
--- a/numba/np/arrayobj.py
+++ b/numba/np/arrayobj.py
@@ -6286,6 +6286,102 @@ def as_strided_impl(x, shape=None, strides=None):
     return as_strided_impl
 
 
+if numpy_version >= (1, 20):
+    @extending.overload(np.lib.stride_tricks.sliding_window_view)
+    def sliding_window_view(x, window_shape, axis=None):
+
+        # Window shape must be given as either an integer or tuple of integers.
+        # We also need to generate buffer tuples we can modify to contain the
+        # final shape and strides (reshape_unchecked does not accept lists).
+        if isinstance(window_shape, types.Integer):
+            shape_buffer = tuple(range(x.ndim + 1))
+            stride_buffer = tuple(range(x.ndim + 1))
+
+            @register_jitable
+            def get_window_shape(window_shape):
+                return (window_shape,)
+
+        elif (isinstance(window_shape, types.UniTuple) and
+              isinstance(window_shape.dtype, types.Integer)):
+            shape_buffer = tuple(range(x.ndim + len(window_shape)))
+            stride_buffer = tuple(range(x.ndim + len(window_shape)))
+
+            @register_jitable
+            def get_window_shape(window_shape):
+                return window_shape
+
+        else:
+            raise errors.TypingError(
+                "window_shape must be an integer or tuple of integers"
+            )
+
+        # Axis must be integer, tuple of integers, or None for all axes.
+        if is_nonelike(axis):
+            @register_jitable
+            def get_axis(window_shape, axis, ndim):
+                return list(range(ndim))
+
+        elif isinstance(axis, types.Integer):
+            @register_jitable
+            def get_axis(window_shape, axis, ndim):
+                return [
+                    normalize_axis("sliding_window_view", "axis", ndim, axis)
+                ]
+
+        elif (isinstance(axis, types.UniTuple) and
+              isinstance(axis.dtype, types.Integer)):
+            @register_jitable
+            def get_axis(window_shape, axis, ndim):
+                return [normalize_axis("sliding_window_view", "axis", ndim, a)
+                        for a in axis]
+
+        else:
+            raise errors.TypingError(
+                "axis must be None, an integer or tuple of integers"
+            )
+
+        def sliding_window_view_impl(x, window_shape, axis=None):
+            window_shape = get_window_shape(window_shape)
+            axis = get_axis(window_shape, axis, x.ndim)
+            if len(window_shape) != len(axis):
+                raise ValueError(
+                    "Must provide matching length window_shape and axis"
+                )
+
+            # Initialise view details with shape and strides of x.
+            out_shape = shape_buffer
+            out_strides = stride_buffer
+            for i in range(x.ndim):
+                out_shape = tuple_setitem(out_shape, i, x.shape[i])
+                out_strides = tuple_setitem(out_strides, i, x.strides[i])
+
+            # Trim the dimensions being windowed and set the window shape and
+            # strides. Note: the same axis can be windowed repeatedly.
+            i = x.ndim
+            for ax, dim in zip(axis, window_shape):
+                if dim < 0:
+                    raise ValueError(
+                        "`window_shape` cannot contain negative values"
+                    )
+                if out_shape[ax] < dim:
+                    raise ValueError(
+                        "window_shape cannot be larger than input array shape"
+                    )
+
+                trimmed = out_shape[ax] - dim + 1
+                out_shape = tuple_setitem(out_shape, ax, trimmed)
+                out_shape = tuple_setitem(out_shape, i, dim)
+                out_strides = tuple_setitem(out_strides, i, x.strides[ax])
+                i += 1
+
+            # The NumPy version calls as_strided, but our implementation of
+            # as_strided is effectively a wrapper for reshape_unchecked.
+            view = reshape_unchecked(x, out_shape, out_strides)
+            return view
+
+        return sliding_window_view_impl
+
+
 @overload(bool)
 def ol_bool(arr):
     if isinstance(arr, types.Array):


Test patch:
diff --git a/numba/tests/test_array_manipulation.py b/numba/tests/test_array_manipulation.py
--- a/numba/tests/test_array_manipulation.py
+++ b/numba/tests/test_array_manipulation.py
@@ -147,6 +147,11 @@ def as_strided2(a):
     return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)
 
 
+@njit
+def sliding_window_view(x, window_shape, axis=None):
+    return np.lib.stride_tricks.sliding_window_view(x, window_shape, axis=axis)
+
+
 def bad_index(arr, arr2d):
     x = arr.x,
     y = arr.y
@@ -516,6 +521,112 @@ def test_as_strided(self):
         self.check_as_strided(as_strided1)
         self.check_as_strided(as_strided2)
 
+    @unittest.skipIf(numpy_version < (1, 20), "requires NumPy 1.20 or newer")
+    def test_sliding_window_view(self):
+        def check(arr, window_shape, axis):
+            # Our version is always writeable (NumPy default is False).
+            expected = np.lib.stride_tricks.sliding_window_view(
+                arr, window_shape, axis, writeable=True
+            )
+            got = sliding_window_view(arr, window_shape, axis)
+            self.assertPreciseEqual(got, expected)
+
+        # 1d array, different ways of specifying the axis.
+        arr1 = np.arange(24)
+        for axis in [None, 0, -1, (0,)]:
+            with self.subTest(f"1d array, axis={axis}"):
+                check(arr1, 5, axis)
+
+        # 2d array, 1d window.
+        arr2 = np.arange(200).reshape(10, 20)
+        for axis in [0, -1]:
+            with self.subTest(f"2d array, axis={axis}"):
+                check(arr2, 5, axis)
+
+        # 2d array, 2d window.
+        for axis in [None, (0, 1), (1, 0), (1, -2)]:
+            with self.subTest(f"2d array, axis={axis}"):
+                check(arr2, (5, 8), axis)
+
+        # 4d array, 2d window.
+        arr4 = np.arange(200).reshape(4, 5, 5, 2)
+        for axis in [(1, 2), (-2, -3)]:
+            with self.subTest(f"4d array, axis={axis}"):
+                check(arr4, (3, 2), axis)
+
+        # Repeated axis.
+        with self.subTest("2d array, repeated axes"):
+            check(arr2, (5, 3, 3), (0, 1, 0))
+
+    @unittest.skipIf(numpy_version < (1, 20), "requires NumPy 1.20 or newer")
+    def test_sliding_window_view_errors(self):
+        def _raises(msg, *args):
+            with self.assertRaises(ValueError) as raises:
+                sliding_window_view(*args)
+            self.assertIn(msg, str(raises.exception))
+
+        def _typing_error(msg, *args):
+            with self.assertRaises(errors.TypingError) as raises:
+                sliding_window_view(*args)
+            self.assertIn(msg, str(raises.exception))
+
+        # Exceptions leak references
+        self.disable_leak_check()
+
+        arr1 = np.arange(24)
+        arr2 = np.arange(200).reshape(10, 20)
+
+        # Window shape cannot be larger than dimension or negative.
+        with self.subTest("1d window shape too large"):
+            _raises("window_shape cannot be larger", arr1, 25, None)
+        with self.subTest("2d window shape too large"):
+            _raises("window_shape cannot be larger", arr2, (4, 21), None)
+        with self.subTest("1d window negative size"):
+            _raises("`window_shape` cannot contain negative", arr1, -1, None)
+        with self.subTest("2d window with a negative size"):
+            _raises("`window_shape` cannot contain negative", arr2, (4, -3), None)
+
+        # window_shape and axis parameters must be compatible.
+        with self.subTest("1d array, 2d window shape"):
+            _raises("matching length window_shape and axis", arr1, (10, 2), None)
+        with self.subTest("2d window shape, only one axis given"):
+            _raises("matching length window_shape and axis", arr2, (10, 2), 1)
+        with self.subTest("1d window shape, 2 axes given"):
+            _raises("matching length window_shape and axis", arr1, 5, (0, 0))
+
+        # Axis values out of bounds.
+        with self.subTest("1d array, second axis"):
+            _raises("Argument axis out of bounds", arr1, 4, 1)
+        with self.subTest("1d array, axis -2"):
+            _raises("Argument axis out of bounds", arr1, 4, -2)
+        with self.subTest("2d array, fourth axis"):
+            _raises("Argument axis out of bounds", arr2, (4, 4), (0, 3))
+        with self.subTest("2d array, axis -3"):
+            _raises("Argument axis out of bounds", arr2, (4, 4), (0, -3))
+
+        # Useful messages for unsupported types.
+        with self.subTest("window_shape=None"):
+            _typing_error(
+                "window_shape must be an integer or tuple of integer", arr1, None
+            )
+        with self.subTest("window_shape=float"):
+            _typing_error(
+                "window_shape must be an integer or tuple of integer", arr1, 3.1
+            )
+        with self.subTest("window_shape=tuple(float)"):
+            _typing_error(
+                "window_shape must be an integer or tuple of integer", arr1, (3.1,)
+            )
+        with self.subTest("axis=float"):
+            _typing_error(
+                "axis must be None, an integer or tuple of integer", arr1, 4, 3.1
+            )
+        with self.subTest("axis=tuple(float)"):
+            _typing_error(
+                "axis must be None, an integer or tuple of integer", arr1, 4, (3.1,)
+            )
+
+
     def test_flatten_array(self, flags=enable_pyobj_flags, layout='C'):
         a = np.arange(9).reshape(3, 3)
         if layout == 'F':