提交 bcc3bd89 authored 作者: ChienliMa's avatar ChienliMa

PerformLinker.make_all() can accepts storage_map

上级 1616cb13
......@@ -545,7 +545,7 @@ def map_storage(fgraph, order, input_storage, output_storage, storage_map = None
input_storage = [[None] for input in fgraph.inputs]
else:
assert len(fgraph.inputs) == len(input_storage)
if storage_map is None:
storage_map = {}
......@@ -759,7 +759,7 @@ class PerformLinker(LocalLinker):
self.no_recycling = no_recycling
return self
def make_all(self, input_storage=None, output_storage=None):
def make_all(self, input_storage=None, output_storage=None, storage_map=None):
"""
Parameters
......@@ -781,7 +781,7 @@ class PerformLinker(LocalLinker):
order = self.schedule(fgraph)
no_recycling = self.no_recycling
input_storage, output_storage, storage_map = map_storage(fgraph, order, input_storage, output_storage)
input_storage, output_storage, storage_map = map_storage(fgraph, order, input_storage, output_storage, storage_map)
compute_map = {}
for k in storage_map:
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论