from __future__ import absolute_import, print_function, division from collections import defaultdict from six import iteritems from theano.gof.graph import list_of_nodes from theano.compat import cmp # {{{ http://code.activestate.com/recipes/578231/ (r1) # Copyright (c) Oren Tirosh 2012 # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. def memodict(f): """ Memoization decorator for a function taking a single argument. """ class memodict(defaultdict): def __missing__(self, key): ret = self[key] = f(key) return ret return memodict().__getitem__ # end of http://code.activestate.com/recipes/578231/ }}} def make_depends(): @memodict def depends(pair): """ Returns True if a depends on b. """ a, b = pair return (any(bout in a.inputs for bout in b.outputs) or any(depends((ainp.owner, b)) for ainp in a.inputs if ainp.owner)) return depends def make_dependence_cmp(): """ Create a comparator to represent the dependence of nodes in a graph. """ depends = make_depends() def dependence(a, b): """ A cmp function for nodes in a graph - does a depend on b? Returns ------- int Positive number if a depends on b, negative number if b depends on a, 0 otherwise. """ if depends((a, b)): return 1 if depends((b, a)): return -1 return 0 return dependence def reverse_dict(d): """ Reverses direction of dependence dict. Notes ----- dict order is not deterministic. As we iterate on the input dict, it makes the output of this function depend on the dict order. So this function output order should be considered as undeterministic. Examples -------- >>> d = {'a': (1, 2), 'b': (2, 3), 'c':()} >>> reverse_dict(d) {1: ('a',), 2: ('a', 'b'), 3: ('b',)} """ result = {} for key in d: for val in d[key]: result[val] = result.get(val, tuple()) + (key, ) return result def _toposort(edges): """ Topological sort algorithm by Kahn [1] - O(nodes + vertices). Parameters ---------- edges A dict of the form {a: {b, c}} where b and c depend on a. Returns ------- L : list An ordered list of nodes that satisfy the dependencies of edges. Closely follows the wikipedia page [2] References ---------- [1] Kahn, Arthur B. (1962), "Topological sorting of large networks", Communications of the ACM [2] http://en.wikipedia.org/wiki/Toposort#Algorithms Examples -------- >>> _toposort({1: {2, 3}, 2: (3, )}) [1, 2, 3] """ incoming_edges = reverse_dict(edges) incoming_edges = dict((k, set(val)) for k, val in iteritems(incoming_edges)) S = set((v for v in edges if v not in incoming_edges)) L = [] while S: n = S.pop() L.append(n) for m in edges.get(n, ()): assert n in incoming_edges[m] incoming_edges[m].remove(n) if not incoming_edges[m]: S.add(m) if any(incoming_edges.get(v, None) for v in edges): raise ValueError("Input has cycles") return L def posort(l, *cmps): """ Partially ordered sort with multiple comparators. Given a list of comparators, orders the elements in l so that the comparators are satisfied as much as possible giving precedence to earlier comparators. Parameters ---------- l An iterable of nodes in a graph. cmps A sequence of comparator functions that describe which nodes should come before which others. Returns ------- list A list of nodes which satisfy the comparators as much as possible. Notes ----- Implemented with _toposort. Examples -------- >>> lower_tens = lambda a, b: a/10 - b/10 # prefer lower numbers div 10 >>> prefer evens = lambda a, b: a%2 - b%2 # prefer even numbers >>> posort(list(range(20)), lower_tens, prefer_evens) [0, 8, 2, 4, 6, 1, 3, 5, 7, 9, 16, 18, 10, 12, 14, 17, 19, 11, 13, 15] """ comes_before = dict((a, set()) for a in l) comes_after = dict((a, set()) for a in l) def add_links(a, b): # b depends on a comes_after[a].add(b) comes_after[a].update(comes_after[b]) for c in comes_before[a]: comes_after[c].update(comes_after[a]) comes_before[b].add(a) comes_before[b].update(comes_before[a]) for c in comes_after[b]: comes_before[c].update(comes_before[b]) def check(): """ Tests for cycles in manufactured edges. """ for a in l: for b in l: assert not(b in comes_after[a] and a in comes_after[b]) for cmp_fn in cmps: for a in l: for b in l: if cmp_fn(a, b) < 0: # a wants to come before b # if this wouldn't cause a cycle and isn't already known if b not in comes_before[a] and b not in comes_after[a]: add_links(a, b) # check() # debug code return _toposort(comes_after) def sort_apply_nodes(inputs, outputs, cmps): """ Order a graph of apply nodes according to a list of comparators. The following example sorts first by dependence of nodes (this is a topological sort) and then by lexicographical ordering (nodes that start with 'E' come before nodes that start with 'I' if there is no dependence. Examples -------- >>> from theano.gof.graph import sort_apply_nodes, dependence >>> from theano.tensor import matrix, dot >>> x = matrix('x') >>> y = dot(x*2, x+1) >>> str_cmp = lambda a, b: cmp(str(a), str(b)) # lexicographical sort >>> sort_apply_nodes([x], [y], cmps=[dependence, str_cmp]) [Elemwise{add,no_inplace}(x, InplaceDimShuffle{x,x}.0), InplaceDimShuffle{x,x}(TensorConstant{2}), Elemwise{mul,no_inplace}(x, InplaceDimShuffle{x,x}.0), InplaceDimShuffle{x,x}(TensorConstant{1}), dot(Elemwise{mul,no_inplace}.0, Elemwise{add,no_inplace}.0)] """ return posort(list_of_nodes(inputs, outputs), *cmps) def sort_schedule_fn(*cmps): """ Make a schedule function from comparators. See Also -------- sort_apply_nodes """ dependence = make_dependence_cmp() cmps = (dependence,) + cmps def schedule(fgraph): """ Order nodes in a FunctionGraph. """ return sort_apply_nodes(fgraph.inputs, fgraph.outputs, cmps) return schedule def key_to_cmp(key): """ comparator function based on "key" function """ def key_cmp(a, b): return cmp(key(a), key(b)) return key_cmp