Merge pull request #2162 from SinaHonari/issue1914

theano/tensor/tests/test_subtensor.py

@@ -1271,6 +1271,26 @@ class TestAdvancedSubtensor(unittest.TestCase):
                   [5, 6, 7],
                   [.5, .3 + 2.1, .15 + 2.1]]), aval
 
+    def test_advanced_indexing(self):
+        # tests advanced indexing in Theano for 2D and 3D tensors
+        rng = numpy.random.RandomState(utt.seed_rng())
+        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)
+        cmd = f1(0, 1) == a[[0, 1], :]
+        self.assertTrue(numpy.all(cmp))
+
+        aa = rng.uniform(size=(4, 2, 3))
+        bb = theano.shared(aa)
+        k = tensor.iscalar()
+        z = bb[[i, j, k], :, i:k]
+        f2 = theano.function([i, j, k], z)
+        cmd = f2(0, 1, 2) == aa[[0, 1, 2], :, 0:2]
+        self.assertTrue(numpy.all(cmp))
+
 
 class TestInferShape(utt.InferShapeTester):
     def test_infer_shape(self):

theano/tensor/var.py

@@ -16,6 +16,12 @@ from theano.tensor.type import TensorType
 from theano.configparser import config
 
 
+def equal_slices(s1, s2):
+    return (s1.start == s2.start and
+            s1.stop == s2.stop and
+            s1.step == s2.step)
+
+
 class AsTensorError(TypeError):
     """Raised when as_tensor_variable isn't able to create a
     TensorVariable.
@@ -375,15 +381,17 @@ class _tensor_py_operators:
 
         if advanced:
             if (axis is not None
-                and all(a == slice(None) for a in args[:axis])
-                and all(a == slice(None) for a in args[axis + 1:])
+                and all(isinstance(a, slice) and
+                        equal_slices(a, slice(None)) for a in args[:axis])
+                and all(isinstance(a, slice) and
+                        equal_slices(a, slice(None)) for a in args[axis + 1:])
                 and isinstance(args[axis], (
                     numpy.ndarray,
                     list,
                     TensorVariable,
                     TensorConstant,
                     theano.tensor.sharedvar.TensorSharedVariable))):
-                return self.take(arg, axis)
+                return self.take(args[axis], axis)
             else:
                 return theano.tensor.subtensor.advanced_subtensor(self, *args)
         else: