1from __future__ import absolute_import, print_function, division
2# Note: this code was initially copied from the 'pyutools' package by its
3# original author, and re-licensed under Theano's license.
4import numpy as np
5
6import theano
7from theano.compile.mode import Mode
8
9
10class MonitorMode(Mode):
11    """
12    `MonitorMode` is a debug mode to easily step through function execution.
13
14    Its default behavior is to behave like the 'FAST_RUN' mode. By providing
15    either a `pre_func` (called before a node is executed) or a `post_func`
16    (called after a node is executed) monitoring function, the user can inspect
17    node behavior.
18
19    A typical use case is to detect the introduction of NaN values in a graph.
20    For an example of such a use case, see doc/tutorial/debug_faq.txt.
21
22    Parameters
23    ----------
24    pre_func
25        A function to call before executing a thunk, with arguments:
26        - the thunk index
27        - the Apply node
28        - the thunk to be called
29    post_func
30        A function to call after executing a thunk, with the same three
31        arguments as `pre_func`.
32    optimizer
33        The optimizer to use. One may use for instance 'fast_compile' to skip
34        optimizations.
35    linker
36        DO NOT USE. This mode uses its own linker. The parameter is needed to
37        allow selecting optimizers to use.
38
39    """
40
41    def __init__(self, pre_func=None, post_func=None,
42                 optimizer='default', linker=None):
43        self.pre_func = pre_func
44        self.post_func = post_func
45        wrap_linker = theano.gof.WrapLinkerMany([theano.gof.OpWiseCLinker()],
46                                                [self.eval])
47        if optimizer == 'default':
48            optimizer = theano.config.optimizer
49        if (linker is not None and
50                not isinstance(linker.mode, MonitorMode)):
51            raise Exception("MonitorMode can only use its own linker! You "
52                            "should not provide one.", linker)
53
54        super(MonitorMode, self).__init__(wrap_linker, optimizer=optimizer)
55
56    def __getstate__(self):
57        lnk, opt = super(MonitorMode, self).__getstate__()
58        return (lnk, opt, self.pre_func, self.post_func)
59
60    def __setstate__(self, state):
61        lnk, opt, pre_func, post_func = state
62        self.pre_func = pre_func
63        self.post_func = post_func
64        super(MonitorMode, self).__setstate__((lnk, opt))
65
66    def eval(self, i, node, fn):
67        """
68        The method that calls the thunk `fn`.
69
70        """
71        if self.pre_func is not None:
72            self.pre_func(i, node, fn)
73        fn()
74        if self.post_func is not None:
75            self.post_func(i, node, fn)
76
77    def clone(self, link_kwargs=None, optimizer="", **kwargs):
78        """
79        Create a new instance of this Mode.
80
81        Keyword arguments can be provided for the linker, but they will be
82        ignored, because MonitorMode needs to use its own linker.
83
84        """
85        if optimizer == "":
86            optimizer = self.provided_optimizer
87        new_mode = type(self)(pre_func=self.pre_func,
88                              post_func=self.post_func,
89                              linker=None,
90                              optimizer=optimizer)
91        return new_mode
92
93
94def detect_nan(i, node, fn):
95    for output in fn.outputs:
96        if (not isinstance(output[0], np.random.RandomState) and
97                np.isnan(output[0]).any()):
98            print('*** NaN detected ***')
99            theano.printing.debugprint(node)
100            print('Inputs : %s' % [input[0] for input in fn.inputs])
101            print('Outputs: %s' % [output[0] for output in fn.outputs])
102            break
103