Merge pull request #6171 from lamblin/fix_gpuadvidx

theano/gpuarray/subtensor.py

@@ -511,40 +511,48 @@ class GpuAdvancedSubtensor(HideC, tensor.AdvancedSubtensor):
 
         x = x.reshape(nshp)
 
-        narrays = 0
         transp = list(range(x.ndim))
+        # number of array-indexed dimensions
         p = 0
-        # ap gives the position of the array in case there is only one.
-        # if there are more than one (narray > 1) it should be ignored.
-        ap = 0
+        # ap represents the axis in the resulting array where the
+        # dimensions indexed by arrays and ints will be inserted.
+        # For instance, if all such dimensions are grouped together,
+        # it corresponds to the index of the first such dimension in the
+        # inital array.  If these dimensions are split (with slices
+        # inbetween), then the resulting dimensions will be moved to the
+        # beginning, and ap will be 0.
+        # If no such dimension has been encountered, ap is None.
+        ap = None
+        # Indicates whether we have already encountered an index (array
+        # or number), and then a slice.
+        slice_after_idx = False
         for k, i in enumerate(list(nidx)):
-            if (isinstance(i, np.ndarray) and
-                    i.ndim != 0):
+            if (isinstance(i, np.ndarray) and i.ndim != 0):
                 transp.remove(k)
                 transp.insert(p, k)
-                ap += k
                 i = nidx.pop(k)
                 nidx.insert(p, i)
                 p += 1
-                narrays += 1
+                if ap is None:
+                    # first non-slice index
+                    ap = k
+                elif slice_after_idx:
+                    # We already encountered at least an array or int, and then
+                    # a slice. Array-indexed axes are not grouped,
+                    # moving to the beginning
+                    ap = 0
             else:
-                if narrays == 0:
-                    try:
-                        i.__index__()
-                        # We shift back the position of the array by the
-                        # number of dimensions that are removed by
-                        # indexing.  If ap is bigger than 0 it means we
-                        # have encountered at least one array.
-                        if ap >= 0:
-                            ap -= 1
-                        # If this index is before the first array then
-                        # we will not move the array back to its
-                        # position.  Mark this by faking that there
-                        # are more than two arrays.  This is crazy
-                        # numpy behaviour so blame them.
-                        narrays = 2
-                    except Exception:
-                        pass
+                try:
+                    i.__index__()
+                    if ap is None:
+                        ap = k
+                    # indices do not break the contiguity of
+                    # array-indexed axes
+                except Exception:
+                    # If we already encountered an array/int index, it
+                    # means future ones will not be grouped.
+                    if ap is not None:
+                        slice_after_idx = True
 
         x = x.transpose(*transp)
 
@@ -552,12 +560,16 @@ class GpuAdvancedSubtensor(HideC, tensor.AdvancedSubtensor):
         x = x.__getitem__(idx_)
 
         if p == 0:
+            assert ap is None
             # The only indexing was through slices and indices.
             # This can happen with symbolic slices for instance.
             # Since no view_map is set, we need to copy the returned value
             out[0] = x.copy()
             return
 
+        # At this point, we should have encountered at least one array
+        assert ap is not None
+
         # flatten the array-indexed dimensions
         shape = ((np.prod(x.shape[0: p]),) +
                  x.shape[p:])
@@ -578,10 +590,9 @@ class GpuAdvancedSubtensor(HideC, tensor.AdvancedSubtensor):
         out_flat_shp = take_idx.shape + x.shape[p:]
         o = out_flat.reshape(out_flat_shp)
 
-        # If there was only one array we need to move the indexed
-        # dimension(s) back to the position of the array, which is
-        # stored in ap.  Note that ap is invalid is narrays != 1.
-        if narrays == 1:
+        if ap != 0:
+            # Put the resulting indexing at the place that NumPy
+            # decided was the right one.
             ntransp = list(range(take_idx.ndim, o.ndim))
             ntransp[ap:ap] = list(range(take_idx.ndim))
             o = o.transpose(*ntransp)

theano/tensor/tests/test_subtensor.py

@@ -1546,7 +1546,8 @@ class TestAdvancedSubtensor(unittest.TestCase):
         typ = tensor.TensorType(self.m.type.dtype, self.ix2.type.broadcastable)
         assert a.type == typ, (a.type, typ)
         f = theano.function([self.m, self.ix1, self.ix12], a,
-                            allow_input_downcast=True)
+                            allow_input_downcast=True,
+                            mode=self.mode)
         aval = f([[.4, .9, .1],
                   [5, 6, 7],
                   [.5, .3, .15]],
@@ -1564,7 +1565,8 @@ class TestAdvancedSubtensor(unittest.TestCase):
 
         assert a.type == self.m.type, (a.type, self.m.type)
         f = theano.function([self.m, self.ix1, self.ix12, inc], [a, g_inc],
-                            allow_input_downcast=True)
+                            allow_input_downcast=True,
+                            mode=self.mode)
         aval, gval = f([[.4, .9, .1],
                         [5, 6, 7],
                         [.5, .3, .15]],
@@ -1584,7 +1586,8 @@ class TestAdvancedSubtensor(unittest.TestCase):
 
         assert a.type == self.m.type, (a.type, self.m.type)
         f = theano.function([self.m, self.ix1, inc], [a, g_inc],
-                            allow_input_downcast=True)
+                            allow_input_downcast=True,
+                            mode=self.mode)
         aval, gval = f([[.4, .9, .1],
                         [5, 6, 7],
                         [.5, .3, .15]],
@@ -1601,7 +1604,8 @@ class TestAdvancedSubtensor(unittest.TestCase):
 
         assert a.type == self.m.type, (a.type, self.m.type)
         f = theano.function([self.m, self.ix1, self.ix2], a,
-                            allow_input_downcast=True)
+                            allow_input_downcast=True,
+                            mode=self.mode)
         aval = f([[.4, .9, .1],
                   [5, 6, 7],
                   [.5, .3, .15]],
@@ -1615,13 +1619,13 @@ class TestAdvancedSubtensor(unittest.TestCase):
 
     def test_advanced_indexing(self):
         # tests advanced indexing in Theano for 2D and 3D tensors
-        rng = np.random.RandomState(utt.seed_rng())
+        rng = np.random.RandomState(utt.fetch_seed())
         a = rng.uniform(size=(3, 3))
         b = theano.shared(a)
         i = tensor.iscalar()
         j = tensor.iscalar()
         z = b[[i, j], :]
-        f1 = theano.function([i, j], z)
+        f1 = theano.function([i, j], z, mode=self.mode)
         cmd = f1(0, 1) == a[[0, 1], :]
         self.assertTrue(cmd.all())
 
@@ -1629,7 +1633,7 @@ class TestAdvancedSubtensor(unittest.TestCase):
         bb = theano.shared(aa)
         k = tensor.iscalar()
         z = bb[[i, j, k], :, i:k]
-        f2 = theano.function([i, j, k], z)
+        f2 = theano.function([i, j, k], z, mode=self.mode)
         cmd = f2(0, 1, 2) == aa[[0, 1, 2], :, 0:2]
         self.assertTrue(cmd.all())
 
@@ -1650,17 +1654,34 @@ class TestAdvancedSubtensor(unittest.TestCase):
         r_idx = np.arange(xx.shape[1])[:, np.newaxis]
         c_idx = np.arange(xx.shape[2])[np.newaxis, :]
 
-        out = X[b_idx, r_idx, c_idx].eval({X: xx})
+        f = theano.function([X], X[b_idx, r_idx, c_idx], mode=self.mode)
+        out = f(xx)
         utt.assert_allclose(out, xx[b_idx, r_idx, c_idx])
 
     def test_adv_sub_slice(self):
         # Reported in https://github.com/Theano/Theano/issues/5898
         var = self.shared(np.zeros([3, 3], dtype=config.floatX))
         slc = tensor.slicetype()
-        f = theano.function([slc], var[slc])
+        f = theano.function([slc], var[slc], mode=self.mode)
         s = slice(1, 3)
         f(s)
 
+    def test_adv_grouped(self):
+        # Reported in https://github.com/Theano/Theano/issues/6152
+        rng = np.random.RandomState(utt.fetch_seed())
+        var_v = rng.rand(3, 63, 4).astype(config.floatX)
+        var = self.shared(var_v)
+        idx1_v = rng.randint(0, 61, size=(5, 4)).astype('int32')
+        idx1 = self.shared(idx1_v)
+        idx2 = tensor.arange(4)
+        out = var[:, idx1, idx2]
+        f = theano.function([], out, mode=self.mode)
+        out_v = f()
+        assert out_v.shape == (3, 5, 4)
+
+        out_np = var_v[:, idx1_v, np.arange(4)]
+        utt.assert_allclose(out_v, out_np)
+
     def test_grad(self):
         ones = np.ones((1, 3), dtype=self.dtype)
         n = self.shared(ones * 5, broadcastable=(True, False))