Skip to content

Commit 070f5e6

Browse files
committed
Less myopic check for protected inputs
1 parent 276ed49 commit 070f5e6

File tree

1 file changed

+15
-1
lines changed

1 file changed

+15
-1
lines changed

pytensor/graph/destroyhandler.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,19 @@ def inplace_candidates(fgraph, inputs, protected_inputs=None):
233233
Inputs Variable that you want to use as inplace destination.
234234
235235
"""
236+
237+
def view_is_protected(view_dict, variable, protected):
238+
"""Check if variable or a memory aliased ancestor of variable is in protected list."""
239+
if variable in protected:
240+
return True
241+
try:
242+
while True:
243+
variable = view_dict[variable]
244+
if variable in protected:
245+
return True
246+
except KeyError:
247+
return False
248+
236249
if protected_inputs is None:
237250
from pytensor.compile.function.types import Supervisor
238251

@@ -244,14 +257,15 @@ def inplace_candidates(fgraph, inputs, protected_inputs=None):
244257
protected_inputs.update(fgraph.outputs)
245258

246259
has_destroyers = fgraph.has_destroyers
260+
view_i = fgraph.destroy_handler.view_i
247261

248262
return [
249263
inp
250264
# Remove duplicates, while preserving order by using dict.fromkeys
251265
for inp in dict.fromkeys(inputs)
252266
if (
253267
not isinstance(inp, Constant)
254-
and inp not in protected_inputs
268+
and not view_is_protected(view_i, inp, protected_inputs)
255269
and not has_destroyers([inp])
256270
)
257271
]

0 commit comments

Comments
 (0)