Skip to content

Commit ad1af2e

Browse files
committed
Delay creation of Scan inner mode
1 parent de75bba commit ad1af2e

File tree

1 file changed

+10
-14
lines changed

1 file changed

+10
-14
lines changed

pytensor/scan/op.py

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@
5858
from pytensor.compile.builders import construct_nominal_fgraph, infer_shape
5959
from pytensor.compile.function.pfunc import pfunc
6060
from pytensor.compile.io import In, Out
61-
from pytensor.compile.mode import Mode, get_default_mode, get_mode
61+
from pytensor.compile.mode import Mode, get_mode
6262
from pytensor.compile.profiling import register_profiler_printer
6363
from pytensor.configdefaults import config
6464
from pytensor.gradient import DisconnectedType, NullType, Rop, grad, grad_undefined
@@ -761,18 +761,7 @@ def __init__(
761761
self.profile = profile
762762
self.allow_gc = allow_gc
763763
self.strict = strict
764-
765-
# Clone mode_instance, altering "allow_gc" for the linker,
766-
# and adding a message if we profile
767-
if self.name:
768-
message = f"{self.name} sub profile"
769-
else:
770-
message = "Scan sub profile"
771-
772-
self.mode = get_default_mode() if mode is None else mode
773-
self.mode_instance = get_mode(self.mode).clone(
774-
link_kwargs=dict(allow_gc=self.allow_gc), message=message
775-
)
764+
self.mode = mode
776765

777766
# build a list of output types for any Apply node using this op.
778767
self.output_types = []
@@ -1445,10 +1434,17 @@ def fn(self):
14451434
elif self.profile:
14461435
profile = self.profile
14471436

1437+
# Clone mode_instance, altering "allow_gc" for the linker,
1438+
# and adding a message if we profile
1439+
mode_instance = get_mode(self.mode).clone(
1440+
link_kwargs=dict(allow_gc=self.allow_gc),
1441+
message=f"{self.name or 'Scan'} sub profile",
1442+
)
1443+
14481444
self._fn = pfunc(
14491445
wrapped_inputs,
14501446
wrapped_outputs,
1451-
mode=self.mode_instance,
1447+
mode=mode_instance,
14521448
accept_inplace=False,
14531449
profile=profile,
14541450
on_unused_input="ignore",

0 commit comments

Comments
 (0)