Skip to content

Commit

Permalink
Temporarily disable inplace for multiple-output Composites
Browse files Browse the repository at this point in the history
  • Loading branch information
Ricardo Vieira committed Oct 17, 2022
1 parent b2914c6 commit 762061b
Showing 1 changed file with 10 additions and 2 deletions.
12 changes: 10 additions & 2 deletions aesara/tensor/rewriting/elemwise.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,14 @@ def print_profile(cls, stream, prof, level=0):
for n in sorted(ndim.keys()):
print(blanc, n, ndim[n], file=stream)

def candidate_input_idxs(self, node):
if isinstance(node.op.scalar_op, aes.Composite) and len(node.outputs) > 1:
# TODO: Implement specialized InplaceCompositeOptimizer with logic
# needed to correctly assign inplace for multi-output Composites
return []
else:
return range(len(node.outputs))

def apply(self, fgraph):
r"""
Expand Down Expand Up @@ -148,7 +156,7 @@ def apply(self, fgraph):

baseline = op.inplace_pattern
candidate_outputs = [
i for i in range(len(node.outputs)) if i not in baseline
i for i in self.candidate_input_idxs(node) if i not in baseline
]
# node inputs that are Constant, already destroyed,
# or fgraph protected inputs and fgraph outputs can't be used as
Expand All @@ -166,7 +174,7 @@ def apply(self, fgraph):
]
else:
baseline = []
candidate_outputs = list(range(len(node.outputs)))
candidate_outputs = self.candidate_input_idxs(node)
# node inputs that are Constant, already destroyed,
# fgraph protected inputs and fgraph outputs can't be used as inplace
# target.
Expand Down

0 comments on commit 762061b

Please sign in to comment.