solve confilcts in test_opt.py

theano/tensor/opt.py

@@ -1894,6 +1894,21 @@ def local_track_shape_i(node):
         replacement = shape_feature.scheduled[node]
         return [shape_feature.shape_of[replacement][node.op.i]]
 
+#TH-2801 opt: subtensor(incsubtensor)
+@register_specialize
+@register_canonicalize
+@gof.local_optimizer([Subtensor])
+def local_subtensor_inc_subtensor(node):
+    if isinstance(node.op, Subtensor):
+        x = node.inputs[0]
+        if not x.owner or not isinstance(x.owner.op, IncSubtensor):
+            return
+        if not x.owner.op.set_instead_of_inc:
+            return
+        if x.owner.inputs[2] == node.inputs[1] and tuple(x.owner.op.idx_list) == tuple(node.op.idx_list):
+            return [x.owner.inputs[1]]
+        else:
+            return
 
 @register_specialize
 @register_canonicalize

theano/tensor/tests/test_opt.py

@@ -1934,6 +1934,44 @@ def test_local_subtensor_remove_broadcastable_index():
     f2(xn)
 
 
+def test_subtensor_inc_subtensor():
+    # basic test
+    x = tensor.matrix('x')
+    i = tensor.iscalar('i')
+    v = tensor.vector('v')
+    y = tensor.set_subtensor(x[i], v)
+    z = y[i]
+    mode = theano.compile.mode.get_default_mode().including('local_subtensor_inc_subtensor')
+    f = theano.function([x, i, v], z, mode=mode)
+    prog = f.maker.fgraph.toposort()
+    assert len(prog) == 1
+    assert isinstance(prog[0].op, DeepCopyOp)
+
+    # complicated test
+    x = tensor.tensor4('x')
+    i1 = tensor.iscalar('i1')
+    i2 = tensor.iscalar('i2')
+    i3 = tensor.iscalar('i3')
+    i4 = tensor.iscalar('i4')
+    v = tensor.tensor3('v')
+    y = tensor.set_subtensor(x[i1, :i2, i3:, ::i4], v)
+    z = y[i1, :i2, i3:, ::i4]
+    mode = theano.compile.mode.get_mode('FAST_COMPILE').including('local_subtensor_inc_subtensor')
+    f = theano.function([x, i1, i2, i3, i4, v], z, mode=mode)
+    prog = f.maker.fgraph.toposort()
+    assert len(prog) == 1
+    assert isinstance(prog[0].op, DeepCopyOp)
+
+    # case not use this optimization
+    z = y[i1, :i3, i2:, ::i4]
+    mode = theano.compile.mode.get_mode('FAST_COMPILE').including('local_subtensor_inc_subtensor')
+    f = theano.function([x, i1, i2, i3, i4, v], z, mode=mode)
+    prog = f.maker.fgraph.toposort()
+    assert len(prog) != 1
+    assert any(isinstance(x.op, tensor.IncSubtensor) for x in prog)
+    assert any(isinstance(x.op, tensor.Subtensor) for x in prog)
+
+
 class test_local_subtensor_make_vector(unittest.TestCase):
     def test_scalar_idx(self):
         x, y, z = tensor.lscalars('xyz')