提交 dae731d1 authored 作者: Ricardo Vieira's avatar Ricardo Vieira 提交者: Ricardo Vieira

Optimize blockwise fallback gufunc function

上级 e39fda37
...@@ -308,17 +308,23 @@ class Blockwise(Op): ...@@ -308,17 +308,23 @@ class Blockwise(Op):
# Wrap core_op perform method in numpy vectorize # Wrap core_op perform method in numpy vectorize
n_outs = len(self.outputs_sig) n_outs = len(self.outputs_sig)
core_node = self._create_dummy_core_node(node.inputs) core_node = self._create_dummy_core_node(node.inputs)
inner_outputs_storage = [[None] for _ in range(n_outs)]
def core_func(*inner_inputs): def core_func(
inner_outputs = [[None] for _ in range(n_outs)] *inner_inputs,
core_node=core_node,
inner_inputs = [np.asarray(inp) for inp in inner_inputs] inner_outputs_storage=inner_outputs_storage,
self.core_op.perform(core_node, inner_inputs, inner_outputs) ):
self.core_op.perform(
core_node,
[np.asarray(inp) for inp in inner_inputs],
inner_outputs_storage,
)
if len(inner_outputs) == 1: if n_outs == 1:
return inner_outputs[0][0] return inner_outputs_storage[0][0]
else: else:
return tuple(r[0] for r in inner_outputs) return tuple(r[0] for r in inner_outputs_storage)
gufunc = np.vectorize(core_func, signature=self.signature) gufunc = np.vectorize(core_func, signature=self.signature)
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论