adding advanced indexing to theano tensors

theano/tensor/tests/test_subtensor.py

@@ -1429,3 +1429,24 @@ class TestInferShape(utt.InferShapeTester):
         self._compile_and_check([admat, advec],
                     [set_subtensor(admat[aivec_val, bivec_val], advec)],
                     [admat_val, advec_val], AdvancedIncSubtensor)
+
+
+def test_advanced_indexing():
+    # 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 = T.iscalar()
+    j = T.iscalar()
+    z = b[[i, j], :]
+    f1 = theano.function([i,j],z)
+    cmd = f1(0,1) == a[[0,1],:]
+    numpy.all(cmp)
+
+    aa = rng.uniform(size=(4,2,3))
+    bb = theano.shared(aa)
+    k = T.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]
+    numpy.all(cmp)

theano/tensor/var.py

@@ -350,6 +350,11 @@ class _tensor_py_operators:
     # argument slice(1, None, None), which is much more desirable.
     # __getslice__ is deprecated in python 2.6 anyway.
 
+    def equal_slices(self, s1, s2):
+        return (s1.start == s2.start and
+                    s1.stop == s2.stop and
+                    s1.step == s2.step)
+
     def __getitem__(self, args):
         if not isinstance(args, tuple):
             args = args,
@@ -375,15 +380,15 @@ 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(self.qual_slices(a, slice(None)) for a in args[:axis])
+                and all(self.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: