Merge pull request #3 from jaberg/inc_advsubtensor

theano/tensor/basic.py

@@ -3491,10 +3491,9 @@ def inc_subtensor(x, y, inplace=False, set_instead_of_inc=False,
     >>> new_r = inc_subtensor(r[10:], 5)
     """
     # retrieve idx_list from x.owner
-    if not isinstance(x.owner.op, Subtensor):
-        raise TypeError('x must be result of a subtensor operation')
+    if isinstance(x.owner.op, Subtensor):
         if tolerate_inplace_aliasing:
-        destroyhandler_tolerate_aliased = [[0,1]]
+            destroyhandler_tolerate_aliased = [[0, 1]]
         else:
             destroyhandler_tolerate_aliased = []
         the_op = IncSubtensor(x.owner.op.idx_list, inplace, set_instead_of_inc,
@@ -3502,6 +3501,18 @@ def inc_subtensor(x, y, inplace=False, set_instead_of_inc=False,
         real_x = x.owner.inputs[0]
         real_idxargs = x.owner.inputs[1:]
         return the_op(real_x, y, *real_idxargs)
+    elif isinstance(x.owner.op, AdvancedSubtensor1):
+        real_x = x.owner.inputs[0]
+        ilist = x.owner.inputs[1]
+        if set_instead_of_inc:
+            the_op = AdvancedIncSubtensor1(inplace, set_instead_of_inc=True)
+        else:
+            the_op = AdvancedIncSubtensor1(inplace, set_instead_of_inc=False)
+        return the_op(real_x, y, ilist)
+    elif isinstance(x.owner.op, AdvancedSubtensor):
+        raise NotImplementedError()
+    else:
+        raise TypeError('x must be result of a subtensor operation')
 
 
 class IncSubtensor(Op):
@@ -4882,33 +4893,36 @@ advanced_subtensor1 = AdvancedSubtensor1()
 
 class AdvancedIncSubtensor1(Op):
     """Increments a subtensor using advanced slicing (list of index)"""
-    def __init__(self, inplace=False):
+    def __init__(self, inplace=False, set_instead_of_inc=False):
         self.inplace = inplace
+        self.set_instead_of_inc = set_instead_of_inc
         if inplace:
             self.destroy_map = {0: [0]}
 
     def __hash__(self):
-        return hash(type(self))
+        return hash((type(self), self.inplace, self.set_instead_of_inc))
+
     def __eq__(self, other):
-        return type(self) == type(other)
+        return (type(self) == type(other)
+                and self.inplace == other.inplace
+                and self.set_instead_of_inc == other.set_instead_of_inc)
 
     def make_node(self, x, y, ilist):
         x_ = as_tensor_variable(x)
         y_ = as_tensor_variable(y)
         ilist_ = as_tensor_variable(ilist)
 
-        assert x_.type.dtype == y_.type.dtype
-        assert x_.type.ndim == y_.type.ndim
-
         if ilist_.type.dtype[:3] not in ('int', 'uin'):
             raise TypeError('index must be integers')
-        if ilist_.type.broadcastable != (False,):
+        if ilist_.type.ndim != 1:
             raise TypeError('index must be vector')
         if x_.type.ndim == 0:
             raise TypeError('cannot index into a scalar')
-        if x_.type.broadcastable[0]:
-            # the caller should have made a copy of x len(ilist) times
-            raise TypeError('cannot index into a broadcastable dimension')
+        if y_.type.ndim > x_.type.ndim:
+            opname = 'increment'
+            raise TypeError('cannot %s x subtensor with ndim=%s'
+            ' by y with ndim=%s to x subtensor with ndim=%s '%(
+                opname, x_.type.ndim, y_.type.ndim ))
 
         return Apply(self, [x_, y_, ilist_], [x_.type()])
 
@@ -4920,8 +4934,21 @@ class AdvancedIncSubtensor1(Op):
             x = x.copy()
         # x[idx] += y don't work if the same index is present many times.
         # It do it only once
+        #  -- Numpy also behaves this way, is it a bug in numpy?
+        if self.set_instead_of_inc:
+            if y.ndim:
+                for (j,i) in enumerate(idx):
+                    x[i] = y[j]
+            else:
+                for i in idx:
+                    x[i] = y
+        else:
+            if y.ndim:
                 for (j,i) in enumerate(idx):
                     x[i] += y[j]
+            else:
+                for i in idx:
+                    x[i] += y
         out[0] = x
 
     def infer_shape(self, node, ishapes):

theano/tensor/tests/test_basic.py

@@ -2376,6 +2376,61 @@ class T_subtensor(unittest.TestCase):
             val = f()
             self.assertTrue(numpy.allclose(val, data[idx].shape))
 
+
+class TestIncSubtensor1(unittest.TestCase):
+    # test inc_subtensor
+    # also tests set_subtensor 
+
+    def setUp(self):
+        self.s = iscalar()
+        self.v = fvector()
+        self.m = dmatrix()
+        self.t = ctensor3()
+
+        self.adv1q = lvector() # advanced 1d query
+
+    def test_cant_adv_idx_into_scalar(self):
+        self.assertRaises(TypeError, lambda : self.s[self.adv1q])
+
+    def test_index_into_vec_w_vec(self):
+        a = self.v[self.adv1q]
+        assert a.type == self.v.type
+
+    def test_1d_set_adv_selection(self):
+        a = set_subtensor(self.v[self.adv1q], self.v[self.adv1q])
+
+        assert a.type == self.v.type
+
+        #TODO: compile a function and verify that the subtensor is removed
+        #      completely, because the whole expression is redundant.
+
+        f = theano.function([self.v, self.adv1q], a, allow_input_downcast=True)
+        aval = f([.4, .9, .1], [1,2])
+        assert numpy.allclose(aval, [.4, 0.9, 0.1])
+
+
+    def test_1d_inc_adv_selection(self):
+        a = inc_subtensor(self.v[self.adv1q], self.v[self.adv1q])
+
+        assert a.type == self.v.type
+        f = theano.function([self.v, self.adv1q], a, allow_input_downcast=True)
+        aval = f([.4, .9, .1], [1,2])
+        assert numpy.allclose(aval, [.4, 1.8, 0.2])
+
+
+    def test_1d_inc_adv_selection_w_broadcasting(self):
+        a = inc_subtensor(self.v[self.adv1q], 3.0)
+
+        assert a.type == self.v.type
+        f = theano.function([self.v, self.adv1q], a, allow_input_downcast=True)
+        aval = f([.4, .9, .1], [1,2])
+        assert numpy.allclose(aval, [.4, 3.9, 3.1])
+
+    def test_assigning_matrix_to_vector_selection(self):
+        self.assertRaises(TypeError,
+                lambda : inc_subtensor(self.v[self.adv1q], fmatrix()))
+
+
 class T_Join_and_Split(unittest.TestCase):
     """
     Split is tested by each verify_grad method.