Skip to content

Commit

Permalink
Optimize blockwise fallback gufunc function
Browse files Browse the repository at this point in the history
  • Loading branch information
ricardoV94 committed Oct 24, 2024
1 parent e39fda3 commit dae731d
Showing 1 changed file with 15 additions and 9 deletions.
24 changes: 15 additions & 9 deletions pytensor/tensor/blockwise.py
Original file line number Diff line number Diff line change
Expand Up @@ -308,17 +308,23 @@ def _create_node_gufunc(self, node) -> None:
# 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,
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,
)

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)

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)

Expand Down

0 comments on commit dae731d

Please sign in to comment.