Skip to content

P
pytensor
提交 f6ab37b5 authored 作者: Maxim Kochurov's avatar Maxim Kochurov 提交者： Maxim Kochurov
浏览文件
操作

clean up pytensor.graph.sched

上级 fc2dc234
显示空白字符变更
内嵌 并排
正在显示 包含 12 行增加 和 284 行删除
pytensor/graph/sched.py deleted 100644 → 0
浏览文件 @ fc2dc234
from collections import defaultdict
from pytensor.graph.basic import list_of_nodes
from pytensor.utils 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 = {k: set(val) for k, val in incoming_edges.items()}
S = {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(nodes, *cmps):
"""
Partially ordered sort with multiple comparators.
Given a list of comparators, orders the elements in `nodes` so that the
comparators are satisfied as much as possible giving precedence to
earlier comparators.
Parameters
----------
nodes
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 = {a: set() for a in nodes}
comes_after = {a: set() for a in nodes}
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 nodes:
for b in nodes:
assert not (b in comes_after[a] and a in comes_after[b])
for cmp_fn in cmps:
for a in nodes:
for b in nodes:
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 pytensor.graph.basic import sort_apply_nodes, dependence
>>> from pytensor.tensor.type import matrix
>>> from pytensor.tensor.math import 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
pytensor/tensor/io.py
浏览文件 @ f6ab37b5
...@@ -2,9 +2,9 @@ import numpy as np ...@@ -2,9 +2,9 @@ import numpy as np
from pytensor.graph.basic import Apply, Constant, Variable from pytensor.graph.basic import Apply, Constant, Variable
from pytensor.graph.op import Op from pytensor.graph.op import Op
from pytensor.graph.sched import key_to_cmp
from pytensor.link.c.type import Generic from pytensor.link.c.type import Generic
from pytensor.tensor.type import tensor from pytensor.tensor.type import tensor
from pytensor.utils import key_to_cmp
class LoadFromDisk(Op): class LoadFromDisk(Op):
......
pytensor/utils.py
浏览文件 @ f6ab37b5
...@@ -111,6 +111,17 @@ def cmp(x, y): ...@@ -111,6 +111,17 @@ def cmp(x, y):
return (x > y) - (x < y) return (x > y) - (x < y)
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
def get_unbound_function(unbound): def get_unbound_function(unbound):
# Op.make_thunk isn't bound, so don't have a __func__ attr. # Op.make_thunk isn't bound, so don't have a __func__ attr.
# But bound method, have a __func__ method that point to the # But bound method, have a __func__ method that point to the
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论