|
58 | 58 | from pytensor.compile.builders import construct_nominal_fgraph, infer_shape |
59 | 59 | from pytensor.compile.function.pfunc import pfunc |
60 | 60 | 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 |
62 | 62 | from pytensor.compile.profiling import register_profiler_printer |
63 | 63 | from pytensor.configdefaults import config |
64 | 64 | from pytensor.gradient import DisconnectedType, NullType, Rop, grad, grad_undefined |
@@ -761,18 +761,7 @@ def __init__( |
761 | 761 | self.profile = profile |
762 | 762 | self.allow_gc = allow_gc |
763 | 763 | 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 |
776 | 765 |
|
777 | 766 | # build a list of output types for any Apply node using this op. |
778 | 767 | self.output_types = [] |
@@ -1445,10 +1434,17 @@ def fn(self): |
1445 | 1434 | elif self.profile: |
1446 | 1435 | profile = self.profile |
1447 | 1436 |
|
| 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 | + |
1448 | 1444 | self._fn = pfunc( |
1449 | 1445 | wrapped_inputs, |
1450 | 1446 | wrapped_outputs, |
1451 | | - mode=self.mode_instance, |
| 1447 | + mode=mode_instance, |
1452 | 1448 | accept_inplace=False, |
1453 | 1449 | profile=profile, |
1454 | 1450 | on_unused_input="ignore", |
|
0 commit comments