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