Repository: numba/numba

Issue:
Add support for `as_strided` from numpy.lib.stride_tricks
With `numba.__version__ = '0.25.0'` trying this code:

```
from numpy.lib.stride_tricks import as_strided
@jit(nopython=True)
def rolling_window(a, window):
    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
    strides = a.strides + (a.strides[-1],)
    return as_strided(a, shape=shape, strides=strides)
```

Leads to:

```
Failed at nopython (nopython frontend)
Untyped global name 'as_strided'
File "<ipython-input-100-5981db08376a>", line 6
```

My understanding is that `as_strided` needs to be added to supported numpy functions for the above code to work. Is that right or there is something that needs to be changed in the code for it to work?



Discussion:
That's correct.  We don't have support for `as_strided`.  I'm not sure if we support tuple concatenation, but we could work around that for the moment.


Resolving patch:
diff --git a/numba/targets/arrayobj.py b/numba/targets/arrayobj.py
--- a/numba/targets/arrayobj.py
+++ b/numba/targets/arrayobj.py
@@ -14,7 +14,7 @@
 
 import numpy as np
 
-from numba import types, cgutils, typing, utils
+from numba import types, cgutils, typing, utils, extending
 from numba.numpy_support import as_dtype, carray, farray
 from numba.numpy_support import version as numpy_version
 from numba.targets.imputils import (lower_builtin, lower_getattr,
@@ -3457,3 +3457,82 @@ def array_to_array(context, builder, fromty, toty, val):
     # Type inference should have prevented illegal array casting.
     assert toty.layout == 'A'
     return val
+
+
+# -----------------------------------------------------------------------------
+# Stride tricks
+
+def reshape_unchecked(a, shape, strides):
+    """
+    An intrinsic returning a derived array with the given shape and strides.
+    """
+    raise NotImplementedError
+
+@extending.type_callable(reshape_unchecked)
+def type_reshape_unchecked(context):
+    def check_shape(shape):
+        return (isinstance(shape, types.BaseTuple) and
+                all(isinstance(v, types.Integer) for v in shape))
+
+    def typer(a, shape, strides):
+        if not isinstance(a, types.Array):
+            return
+        if not check_shape(shape) or not check_shape(strides):
+            return
+        if len(shape) != len(strides):
+            return
+        return a.copy(ndim=len(shape), layout='A')
+
+    return typer
+
+@lower_builtin(reshape_unchecked, types.Array, types.BaseTuple, types.BaseTuple)
+def impl_shape_unchecked(context, builder, sig, args):
+    aryty = sig.args[0]
+    retty = sig.return_type
+
+    ary = make_array(aryty)(context, builder, args[0])
+    out = make_array(retty)(context, builder)
+    shape = cgutils.unpack_tuple(builder, args[1])
+    strides = cgutils.unpack_tuple(builder, args[2])
+
+    populate_array(out,
+                   data=ary.data,
+                   shape=shape,
+                   strides=strides,
+                   itemsize=ary.itemsize,
+                   meminfo=ary.meminfo,
+                   )
+
+    res = out._getvalue()
+    return impl_ret_borrowed(context, builder, retty, res)
+
+
+@extending.overload(np.lib.stride_tricks.as_strided)
+def as_strided(x, shape=None, strides=None):
+    from numba import jit
+
+    if shape in (None, types.none):
+        @jit(nopython=True)
+        def get_shape(x, shape):
+            return x.shape
+    else:
+        @jit(nopython=True)
+        def get_shape(x, shape):
+            return shape
+
+    if strides in (None, types.none):
+        # When *strides* is not passed, as_strided() does a non-size-checking
+        # reshape(), possibly changing the original strides.  This is too
+        # cumbersome to support right now, and a Web search shows all example
+        # use cases of as_strided() pass explicit *strides*.
+        raise NotImplementedError("as_strided() strides argument is mandatory")
+    else:
+        @jit(nopython=True)
+        def get_strides(x, strides):
+            return strides
+
+    def as_strided_impl(x, shape=None, strides=None):
+        x = reshape_unchecked(x, get_shape(x, shape), get_strides(x, strides))
+        return x
+
+    return as_strided_impl
diff --git a/numba/targets/base.py b/numba/targets/base.py
--- a/numba/targets/base.py
+++ b/numba/targets/base.py
@@ -147,6 +147,7 @@ def append(self, value, sig):
         self.versions.append((sig, value))
         self._cache.clear()
 
+
 @utils.runonce
 def _load_global_helpers():
     """


Test patch:
diff --git a/numba/tests/support.py b/numba/tests/support.py
--- a/numba/tests/support.py
+++ b/numba/tests/support.py
@@ -55,12 +55,17 @@ def compile(self, func, args, return_type=None, flags=DEFAULT_FLAGS):
         Compile the function or retrieve an already compiled result
         from the cache.
         """
+        from numba.targets.registry import cpu_target
+
         cache_key = (func, args, return_type, flags)
         try:
             cr = self.cr_cache[cache_key]
         except KeyError:
-            cr = compile_extra(self.typingctx, self.targetctx, func,
-                               args, return_type, flags, locals={})
+            # Register the contexts in case for nested @jit or @overload calls
+            # (same as compile_isolated())
+            with cpu_target.nested_context(self.typingctx, self.targetctx):
+                cr = compile_extra(self.typingctx, self.targetctx, func,
+                                   args, return_type, flags, locals={})
             self.cr_cache[cache_key] = cr
         return cr
 
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
@@ -44,6 +44,19 @@ def squeeze_array(a):
     return a.squeeze()
 
 
+def as_strided1(a):
+    # as_strided() with implicit shape
+    strides = (a.strides[0] // 2,) + a.strides[1:]
+    return np.lib.stride_tricks.as_strided(a, strides=strides)
+
+def as_strided2(a):
+    # Rolling window example as in https://github.com/numba/numba/issues/1884
+    window = 3
+    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
+    strides = a.strides + (a.strides[-1],)
+    return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)
+
+
 def add_axis1(a):
     return np.expand_dims(a, axis=0)
 
@@ -147,6 +160,24 @@ def check_err_multiple_negative(arr, shape):
         # Exceptions leak references
         self.disable_leak_check()
 
+    def check_as_strided(self, pyfunc):
+        def run(arr):
+            cres = self.ccache.compile(pyfunc, (typeof(arr),))
+            return cres.entry_point(arr)
+        def check(arr):
+            expected = pyfunc(arr)
+            got = run(arr)
+            self.assertPreciseEqual(got, expected)
+
+        arr = np.arange(24)
+        check(arr)
+        check(arr.reshape((6, 4)))
+        check(arr.reshape((4, 1, 6)))
+
+    def test_as_strided(self):
+        self.check_as_strided(as_strided1)
+        self.check_as_strided(as_strided2)
+
     def test_flatten_array(self, flags=enable_pyobj_flags, layout='C'):
         a = np.arange(9).reshape(3, 3)
         if layout == 'F':