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

Optimize blockwise fallback gufunc function

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