|
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