Merge pull request #3991 from vmichals/keep_stack_trace_in_opt

theano/tensor/nnet/conv3d2d.py

@@ -3,6 +3,7 @@ from theano.gradient import DisconnectedType
 from theano.gof import Op, Apply, TopoOptimizer
 from theano import tensor
 import theano.sandbox.cuda as cuda
+from theano.tensor.opt import copy_stack_trace
 
 
 def get_diagonal_subtensor_view(x, i0, i1):
@@ -328,7 +329,11 @@ def make_gpu_optimizer(op, to_gpu):
                 new_inp = list(node.inputs)
                 for idx in to_gpu:
                     new_inp[idx] = cuda.gpu_from_host(new_inp[idx])
-                return [cuda.host_from_gpu(op()(*new_inp))]
+                result_node = op()(*new_inp)
+                copy_stack_trace(node.outputs[0], result_node)
+                transfer_node = cuda.host_from_gpu(result_node)
+                copy_stack_trace(node.outputs[0], transfer_node)
+                return [transfer_node]
         if node.op == cuda.gpu_from_host:
             # gpu_from_host(op) -> op(gpu_from_host)
             host_input = node.inputs[0]
@@ -338,7 +343,9 @@ def make_gpu_optimizer(op, to_gpu):
                 new_inp = list(op_node.inputs)
                 for idx in to_gpu:
                     new_inp[idx] = cuda.gpu_from_host(new_inp[idx])
-                return [op()(*new_inp)]
+                new_node = op()(*new_inp)
+                copy_stack_trace(host_input, new_node)
+                return [new_node]
         return False
     local_to_gpu.__name__ = "local_to_gpu_" + op.__name__
     cuda.opt.register_opt()(local_to_gpu)
@@ -355,6 +362,7 @@ def local_inplace_DiagonalSubtensor(node):
             not node.op.inplace):
         new_op = node.op.__class__(inplace=True)
         new_node = new_op(*node.inputs)
+        copy_stack_trace(node.outputs[0], new_node)
         return [new_node]
     return False
 theano.compile.optdb.register(

theano/tensor/nnet/nnet.py

@@ -752,7 +752,8 @@ def local_logsoftmax(node):
         inVars = node.inputs[0].owner.inputs[0]
         new_op = LogSoftmax()
         ret = new_op(inVars)
-        ret.tag.values_eq_approx = values_eq_approx_remove_inf
+        ret .tag.values_eq_approx = values_eq_approx_remove_inf
+        copy_stack_trace([node.inputs[0], node.outputs[0]], ret)
         return [ret]
 
 
@@ -785,9 +786,9 @@ def local_logsoftmax_grad(node):
         grads = node.inputs[0].owner.inputs[0]
         if grads.broadcastable[1] and not sm.broadcastable[1]:
             grads = tensor.alloc(grads, grads.shape[0], sm.shape[1])
-
         ret = grads - tensor.sum(grads, axis=1, keepdims=True) * sm
         ret.tag.values_eq_approx = values_eq_approx_remove_nan
+        copy_stack_trace(node.outputs[0], ret)
         return [ret]
 
 

theano/tensor/nnet/opt.py

@@ -36,6 +36,7 @@ def local_inplace_sparse_block_gemv(node):
     """
     if isinstance(node.op, SparseBlockGemv) and not node.op.inplace:
         new_node = sparse_block_gemv_inplace(*node.inputs)
+        copy_stack_trace(node.outputs[0], new_node)
         return [new_node]
     return False
 compile.optdb.register('local_inplace_sparse_block_gemv',
@@ -52,6 +53,7 @@ def local_inplace_sparse_block_outer(node):
     """
     if isinstance(node.op, SparseBlockOuter) and not node.op.inplace:
         new_node = sparse_block_outer_inplace(*node.inputs)
+        copy_stack_trace(node.outputs[0], new_node)
         return [new_node]
     return False
 compile.optdb.register('local_inplace_sparse_block_outer',

theano/tensor/nnet/sigm.py

@@ -18,13 +18,14 @@ from theano.printing import pprint
 from theano.tensor import basic as tensor
 from theano.tensor import elemwise, opt, NotScalarConstantError
 from theano.tensor.type import values_eq_approx_remove_inf
-
+from theano.tensor.opt import copy_stack_trace
 
 ############
 #
 # SCALAR OPS
 #
 
+
 class ScalarSigmoid(scalar.UnaryScalarOp):
     """
     This is just speed opt. Not for stability.
@@ -262,6 +263,7 @@ def local_ultra_fast_sigmoid(node):
     if (isinstance(node.op, tensor.Elemwise) and
             node.op.scalar_op == scalar_sigmoid):
         out = ultra_fast_sigmoid(node.inputs[0])
+        copy_stack_trace(node.outputs[0], out)
 
         def values_eq_approx_remove_low_prec(a, b):
             # atol is found by trial/error.
@@ -301,6 +303,7 @@ def local_hard_sigmoid(node):
     if (isinstance(node.op, tensor.Elemwise) and
             node.op.scalar_op == scalar_sigmoid):
         out = hard_sigmoid(node.inputs[0])
+        copy_stack_trace(node.outputs[0], out)
 
         def values_eq_approx_remove_low_prec(a, b):
             # atol is found by trial/error.
@@ -925,7 +928,10 @@ def local_sigm_times_exp(node):
     # get rid of them.
     mul_tree = simplify_mul(mul_tree)
     # Recompute final output based on the updated tree.
-    return [compute_mul(mul_tree)]
+    out = compute_mul(mul_tree)
+    # keep the stack trace
+    copy_stack_trace(node.outputs[0], out)
+    return [out]
 
 
 @opt.register_stabilize
@@ -946,10 +952,17 @@ def local_inv_1_plus_exp(node):
             if len(nonconsts) == 1:
                 if nonconsts[0].owner and nonconsts[0].owner.op == tensor.exp:
                     if scalars and numpy.allclose(numpy.sum(scalars), 1):
-                        return opt._fill_chain(
+                        out = opt._fill_chain(
                             sigmoid(
                                 tensor.neg(nonconsts[0].owner.inputs[0])),
                             scalar_inputs)
+                        # keep combined stack traces of
+                        #     exp(x):           nonconsts[0],
+                        #     1 + exp(x):       inv_arg,
+                        #     1 / (1 + exp(x)): node.outputs[0]
+                        copy_stack_trace(
+                            [nonconsts[0], inv_arg, node.outputs[0]], out)
+                        return out
 
 # Registration is below, and conditional.
 
@@ -970,7 +983,9 @@ def local_1msigmoid(node):
             except Exception:
                 return
             if numpy.allclose(numpy.sum(val_l), 1):
-                return [sigmoid(-sub_r.owner.inputs[0])]
+                out = sigmoid(-sub_r.owner.inputs[0])
+                copy_stack_trace([sub_r, node.outputs[0]], out)
+                return [out]
 
 register_local_1msigmoid = False
 # This is False because the Stabilize pattern above

theano/tensor/nnet/tests/test_conv3d2d.py

@@ -73,6 +73,10 @@ def pyconv3d(signals, filters):
                 r_i += o_i[Tf2:o_i_sh0-Tf2, Hf2:-Hf2, Wf2:-Wf2]
     return rval
 
+def check_diagonal_subtensor_view_traces(fn):
+    for apply_node in fn.maker.fgraph.apply_nodes:
+        if isinstance(apply_node.op, (DiagonalSubtensor, IncDiagonalSubtensor)):
+            assert hasattr(apply_node.outputs[0].tag, 'trace')
 
 def test_conv3d(mode=mode_without_gpu, shared=theano.tensor._shared):
     if ndimage is None:
@@ -100,6 +104,7 @@ def test_conv3d(mode=mode_without_gpu, shared=theano.tensor._shared):
                                 updates={s_output: out},
                                 mode=mode)
 
+    check_diagonal_subtensor_view_traces(newconv3d)
     t0 = time.time()
     newconv3d()
     print(time.time() - t0)
@@ -110,6 +115,7 @@ def test_conv3d(mode=mode_without_gpu, shared=theano.tensor._shared):
                                           (s_signals, gsignals)],
                                  mode=mode,
                                  name='grad')
+    check_diagonal_subtensor_view_traces(gnewconv3d)
 
     t0 = time.time()
     gnewconv3d()
@@ -144,6 +150,7 @@ def test_conv3d(mode=mode_without_gpu, shared=theano.tensor._shared):
     newconv3d = theano.function([], [],
                                 updates={s_output: out},
                                 mode=mode)
+    check_diagonal_subtensor_view_traces(newconv3d)
 
     t0 = time.time()
     newconv3d()
@@ -155,6 +162,7 @@ def test_conv3d(mode=mode_without_gpu, shared=theano.tensor._shared):
                                           (s_signals, gsignals)],
                                  mode=mode,
                                  name='grad')
+    check_diagonal_subtensor_view_traces(gnewconv3d)
 
     t0 = time.time()
     gnewconv3d()

theano/tensor/nnet/tests/test_opt.py

@@ -13,6 +13,7 @@ def test_blocksparse_inplace_gemv_opt():
     o = sparse_block_dot(W, h, iIdx, b, oIdx)
 
     f = theano.function([W, h, iIdx, b, oIdx], o)
+    assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
 
     if theano.config.mode == "FAST_COMPILE":
         assert not f.maker.fgraph.toposort()[-1].op.inplace
@@ -33,6 +34,7 @@ def test_blocksparse_inplace_outer_opt():
 
     f = theano.function([W, h, iIdx, b, oIdx],
                         [o, tensor.grad(o.sum(), wrt=W)])
+    assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
 
     if theano.config.mode == "FAST_COMPILE":
         assert not f.maker.fgraph.toposort()[-1].op.inplace

theano/tensor/nnet/tests/test_sigm.py

@@ -126,32 +126,40 @@ class T_sigmoid_opts(unittest.TestCase):
 
             # tests inv_1_plus_exp
             f = theano.function([x], T.fill(x, 1.0) / (1 + T.exp(-x)), mode=m)
+            assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
             assert [node.op for node in f.maker.fgraph.toposort()] == [sigmoid]
             f(data)
             f = theano.function([x], T.fill(x, 1.0) / (2 + T.exp(-x)), mode=m)
+            assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
             assert [node.op for node in f.maker.fgraph.toposort()] != [sigmoid]
             f(data)
             f = theano.function([x], T.fill(x, 1.0) / (1 - T.exp(-x)), mode=m)
+            assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
             assert [node.op for node in f.maker.fgraph.toposort()] != [sigmoid]
             f(data)
             f = theano.function([x], T.fill(x, 1.1) / (1 + T.exp(-x)), mode=m)
+            assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
             assert [node.op for node in f.maker.fgraph.toposort()] != [sigmoid]
             f(data)
 
             # tests inv_1_plus_exp with neg
             f = theano.function([x], T.fill(x, -1.0) / (1 + T.exp(-x)), mode=m)
+            assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
             assert [node.op for node in f.maker.fgraph.toposort()] == [sigmoid,
                     theano.tensor.inplace.neg_inplace]
             f(data)
             f = theano.function([x], T.fill(x, -1.0) / (1 - T.exp(-x)), mode=m)
+            assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
             assert [node.op for node in f.maker.fgraph.toposort()] != [sigmoid,
                     theano.tensor.inplace.neg_inplace]
             f(data)
             f = theano.function([x], T.fill(x, -1.0) / (2 + T.exp(-x)), mode=m)
+            assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
             assert [node.op for node in f.maker.fgraph.toposort()] != [sigmoid,
                     theano.tensor.inplace.neg_inplace]
             f(data)
             f = theano.function([x], T.fill(x, -1.1) / (1 + T.exp(-x)), mode=m)
+            assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
             assert [node.op for node in f.maker.fgraph.toposort()] != [sigmoid,
                     theano.tensor.inplace.neg_inplace]
             f(data)
@@ -162,31 +170,37 @@ class T_sigmoid_opts(unittest.TestCase):
             # = - (sigm(x) * sigm(x))
             f = theano.function([x], (T.fill(x, -1.0) * T.exp(x)) /
                                 ((1 + T.exp(x)) * (1 + T.exp(-x))), mode=m)
+            assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
             assert [node.op for node in f.maker.fgraph.toposort()] == [sigmoid,
                     T.mul]
             f(data)
             f = theano.function([x], (T.fill(x, -1.1) * T.exp(x)) /
                                 ((1 + T.exp(x)) * (1 + T.exp(-x))), mode=m)
+            assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
             assert [node.op for node in f.maker.fgraph.toposort()] != [sigmoid,
                     T.mul, theano.tensor.inplace.neg_inplace]
             f(data)
             f = theano.function([x], (T.fill(x, -1.0) * T.exp(x)) /
                                 ((2 + T.exp(x)) * (1 + T.exp(-x))), mode=m)
+            assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
             assert [node.op for node in f.maker.fgraph.toposort()] != [sigmoid,
                     T.mul, theano.tensor.inplace.neg_inplace]
             f(data)
             f = theano.function([x], (T.fill(x, -1.0) * T.exp(x)) /
                                 ((1 + T.exp(x)) * (2 + T.exp(-x))), mode=m)
+            assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
             assert [node.op for node in f.maker.fgraph.toposort()] != [sigmoid,
                     T.mul, theano.tensor.inplace.neg_inplace]
             f(data)
             f = theano.function([x], (T.fill(x, -1.0) * T.exp(x)) /
                                 ((1 + T.exp(x)) * (1 + T.exp(x))), mode=m)
+            assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
             assert [node.op for node in f.maker.fgraph.toposort()] != [sigmoid,
                     T.mul, theano.tensor.inplace.neg_inplace]
             f(data)
             f = theano.function([x], (T.fill(x, -1.0) * T.exp(x)) /
                                 ((1 + T.exp(x)) * (2 + T.exp(-x))), mode=m)
+            assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
             assert [node.op for node in f.maker.fgraph.toposort()] != [sigmoid,
                     T.mul, theano.tensor.inplace.neg_inplace]
             f(data)
@@ -204,11 +218,13 @@ class T_sigmoid_opts(unittest.TestCase):
 
         # tests exp_over_1_plus_exp
         f = theano.function([x], 1 - T.exp(x) / (1 + T.exp(x)), mode=m)
+        assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
         assert [node.op for node in f.maker.fgraph.toposort()] == [
             tensor.neg, sigmoid_inplace]
 
         # tests inv_1_plus_exp
         f = theano.function([x], 1 - T.fill(x, 1.0) / (1 + T.exp(-x)), mode=m)
+        assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
         assert [node.op for node in f.maker.fgraph.toposort()] == [tensor.neg,
                 sigmoid_inplace]
 
@@ -225,12 +241,15 @@ class T_sigmoid_opts(unittest.TestCase):
         x, y = tensor.vectors('x', 'y')
 
         f = theano.function([x], sigmoid(-x) * tensor.exp(x), mode=m)
+        assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
         match(f, [sigmoid])
 
         f = theano.function([x], sigmoid(x) * tensor.exp(-x), mode=m)
+        assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
         match(f, [tensor.neg, sigmoid])
 
         f = theano.function([x], -(-(-(sigmoid(x)))) * tensor.exp(-x), mode=m)
+        assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
         match(f, [tensor.neg, sigmoid, tensor.neg])
 
         f = theano.function(
@@ -238,6 +257,7 @@ class T_sigmoid_opts(unittest.TestCase):
                 (sigmoid(x) * sigmoid(-y) * -tensor.exp(-x) *
                  tensor.exp(x * y) * tensor.exp(y)),
                 mode=m)
+        assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
         match(f, [sigmoid, tensor.mul, tensor.neg, tensor.exp, sigmoid,
                   tensor.mul])
 
@@ -298,6 +318,7 @@ class T_sigmoid_opts(unittest.TestCase):
         mode = self.get_mode()
         if not isinstance(mode, theano.compile.DebugMode):
             f = theano.function([x, lr], ux, mode=mode)
+            assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
             ux_v = f([[50]], 0.1)
             assert not numpy.isnan(ux_v)
 
@@ -307,12 +328,14 @@ class T_sigmoid_opts(unittest.TestCase):
 
         mode = self.get_mode('local_ultra_fast_sigmoid')
         f = theano.function([x], s, mode=mode)
+        assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
         topo = f.maker.fgraph.toposort()
         assert len(topo) == 1
         assert topo[0].op == sigmoid
 
         mode = self.get_mode().including('local_ultra_fast_sigmoid')
         f = theano.function([x], s, mode=mode)
+        assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
         topo = f.maker.fgraph.toposort()
         assert topo[0].op == ultra_fast_sigmoid
         assert len(topo) == 1
@@ -324,12 +347,14 @@ class T_sigmoid_opts(unittest.TestCase):
 
         mode = self.get_mode('local_hard_sigmoid')
         f = theano.function([x], s, mode=mode)
+        assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
         topo = f.maker.fgraph.toposort()
         assert topo[0].op == sigmoid
         assert len(topo) == 1
 
         mode = self.get_mode().including('local_hard_sigmoid')
         f = theano.function([x], s, mode=mode)
+        assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
         topo = f.maker.fgraph.toposort()
         assert len(topo) > 1
         assert not any([n.op == sigmoid for n in topo])
@@ -352,6 +377,7 @@ class T_softplus_opts(unittest.TestCase):
 
         out = T.log(sigmoid(x))
         f = theano.function([x], out, mode=self.m)
+        assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
         topo = f.maker.fgraph.toposort()
         assert len(topo) == 3
         assert isinstance(topo[0].op.scalar_op, theano.scalar.Neg)
@@ -375,6 +401,7 @@ class T_softplus_opts(unittest.TestCase):
         # Same test with a flatten
         out = T.log(1 - T.flatten(sigmoid(x)))
         f = theano.function([x], out, mode=self.m)
+        assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
         topo = f.maker.fgraph.toposort()
         assert len(topo) == 3
         assert tensor.is_flat(topo[0].outputs[0])
@@ -403,6 +430,7 @@ class T_softplus_opts(unittest.TestCase):
 
         out = T.log(1 + T.exp(x))
         f = theano.function([x], out, mode=self.m)
+        assert hasattr(f.maker.fgraph.outputs[0].tag, 'trace')
         topo = f.maker.fgraph.toposort()
         assert len(topo) == 1
         assert isinstance(topo[0].op.scalar_op,