Allow single dimension slices in mgrid and ogrid

bugbakery · Aug 24, 2020 · 7f1ab2a · 7f1ab2a
1 parent ea65b92
commit 7f1ab2a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/tests/tensor/test_basic.py b/tests/tensor/test_basic.py
@@ -6729,10 +6729,12 @@ def setup_method(self):
 
     def test_mgrid_numpy_equiv(self):
         nmgrid = (
+            [np.mgrid[0:1:0.1]],
             np.mgrid[0:1:0.1, 1:10:1.0, 10:100:10.0],
             np.mgrid[0:2:1, 1:10:1, 10:100:10],
         )
         tmgrid = (
+            [mgrid[0:1:0.1]],
             mgrid[0:1:0.1, 1:10:1.0, 10:100:10.0],
             mgrid[0:2:1, 1:10:1, 10:100:10],
         )
@@ -6742,10 +6744,12 @@ def test_mgrid_numpy_equiv(self):
 
     def test_ogrid_numpy_equiv(self):
         nogrid = (
+            [np.ogrid[0:1:0.1]],
             np.ogrid[0:1:0.1, 1:10:1.0, 10:100:10.0],
             np.ogrid[0:2:1, 1:10:1, 10:100:10],
         )
         togrid = (
+            [ogrid[0:1:0.1]],
             ogrid[0:1:0.1, 1:10:1.0, 10:100:10.0],
             ogrid[0:2:1, 1:10:1, 10:100:10],
         )

diff --git a/theano/tensor/basic.py b/theano/tensor/basic.py
@@ -5937,6 +5937,10 @@ def __init__(self, sparse=False):
 
     def __getitem__(self, *args):
 
+        if isinstance(args[0], slice):
+            sl = args[0]
+            return arange(sl.start or 0, sl.stop, sl.step or 1)
+
         ndim = len(args[0])
         for sl in args[0]:
             if isinstance(sl.step, python_complex):