Merge pull request #5785 from Faruk-Ahmed/inc_to_set_zeros

theano/tensor/opt.py

 from __future__ import absolute_import, print_function, division
-"""
-Tensor optimizations addressing the ops in basic.py.
+""" Tensor optimizations addressing the ops in basic.py.
 """
 # TODO: intelligent merge for mul/add
 # TODO: 0*x -> 0
@@ -3411,6 +3410,24 @@ def local_incsubtensor_of_zeros(node):
             return
 
 
+@register_canonicalize
+@register_specialize
+@gof.local_optimizer([IncSubtensor])
+def local_incsubtensor_of_zeros_to_setsubtensor(node):
+    """
+    IncSubtensor(zeros, x, ...) -> SetSubtensor(zeros, x, ...)
+    """
+    if (isinstance(node.op, (IncSubtensor)) and not node.op.set_instead_of_inc):
+        x = node.inputs[0]
+
+        if isinstance(x, T.Constant) and not np.any(x.data):
+            return [IncSubtensor(node.op.idx_list,
+                                 node.op.inplace,
+                                 set_instead_of_inc=True,
+                                 destroyhandler_tolerate_aliased=node.op.destroyhandler_tolerate_aliased,
+                                 )(*node.inputs)]
+
+
 @register_canonicalize('local_setsubtensor_of_allocs')
 @register_stabilize('local_setsubtensor_of_allocs')
 @gof.local_optimizer([IncSubtensor])

theano/tensor/subtensor.py

@@ -573,8 +573,9 @@ class Subtensor(Op):
         else:
             # For best optimization, we let this as an inc.
             # This allow the opt local_IncSubtensor_serialize to apply first.
-            # We need to implement an optimization that will convert this to a
-            # set subtensor.
+            # We have an optimization that will convert this to a
+            # set subtensor here at:
+            # theano/tensor/opt.py:local_incsubtensor_of_zeros_to_setsubtensor()
             first = IncSubtensor(self.idx_list)(x.zeros_like(),
                                                 gz, *rest)
         return ([first] + [DisconnectedType()()] * len(rest))

theano/tensor/tests/test_opt.py

@@ -3001,6 +3001,37 @@ class Test_alloc_zero(unittest.TestCase):
         assert np.all([not isinstance(n.op, tensor.IncSubtensor)
                        for n in f.maker.fgraph.toposort()])
 
+    def test_incsubtensor_x_zeros(self):
+        x = tensor.constant(np.asarray(np.zeros((4, 4)),
+                                       dtype=config.floatX))
+        y = tensor.matrix()
+        z = tensor.inc_subtensor(x[:4], y)
+        f = theano.function([y], z)
+        inc_nodes = [n for n in f.maker.fgraph.toposort()
+                     if isinstance(n.op, tensor.IncSubtensor)]
+
+        assert(len(inc_nodes) == 1)
+        node_is_set_instead_of_inc = inc_nodes[0].op.set_instead_of_inc
+        mode = theano.config.mode
+        assert((mode != "FAST_COMPILE" and node_is_set_instead_of_inc) or
+               (mode == "FAST_COMPILE" and not node_is_set_instead_of_inc))
+        test_X = np.random.random((4, 4)).astype(config.floatX)
+        utt.assert_allclose(f(test_X), test_X)
+
+        # also check the flag doesn't get set if first input is not zeros:
+        not_all_zeros = np.zeros((4, 4))
+        not_all_zeros[1, 0] = 0.001
+        x = tensor.constant(np.asarray(not_all_zeros, dtype=config.floatX))
+        y = tensor.matrix()
+        z = tensor.inc_subtensor(x[:4], y)
+        f = theano.function([y], z)
+        inc_nodes = [n for n in f.maker.fgraph.toposort()
+                     if isinstance(n.op, tensor.IncSubtensor)]
+        assert(len(inc_nodes) == 1)
+        assert(inc_nodes[0].op.set_instead_of_inc is False)
+        test_X = np.random.random((4, 4)).astype(config.floatX)
+        utt.assert_allclose(f(test_X), test_X + not_all_zeros)
+
     def test_advancedincsubtensor1_allocs0(self):
         x = tensor.matrix()
         y = tensor.matrix()