提交 5a9fcbce authored 作者: Arnaud Bergeron's avatar Arnaud Bergeron

Don't dig through nodes with multiple clients when searching for a node in the merges.

上级 56ba1c88
...@@ -32,10 +32,12 @@ def grab_cpu_scalar(v, nd): ...@@ -32,10 +32,12 @@ def grab_cpu_scalar(v, nd):
return v.dimshuffle(()) return v.dimshuffle(())
def find_node(v, cls): def find_node(v, cls, ignore_clients=False):
# This digs through possibly redundant transfers to for the node # This digs through possibly redundant transfers to for the node
# that has the op class specified. # that has the op class specified. If ignore_clients is False (the
if v.owner is not None: # default) it will only dig through nodes that have a single
# client.
if v.owner is not None and (ignore_clients or v.clients == 1):
if isinstance(v.owner.op, cls): if isinstance(v.owner.op, cls):
return v.owner return v.owner
elif (isinstance(v.owner.op, GpuFromHost) and elif (isinstance(v.owner.op, GpuFromHost) and
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论