Skip to content

Commit e5b81ff

Browse files
committed
Avoid calling tipyify repeatedly on the same variables
1 parent 336ca3a commit e5b81ff

File tree

1 file changed

+12
-5
lines changed

1 file changed

+12
-5
lines changed

pytensor/link/utils.py

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -731,6 +731,8 @@ def fgraph_to_python(
731731
if global_env is None:
732732
global_env = {}
733733

734+
tipifiyed_vars = set()
735+
734736
body_assigns = []
735737
for node in order:
736738
compiled_func = op_conversion_fn(
@@ -742,17 +744,22 @@ def fgraph_to_python(
742744
global_env[local_compiled_func_name] = compiled_func
743745

744746
node_input_names = []
745-
for i in node.inputs:
746-
local_input_name = unique_name(i)
747+
for inp in node.inputs:
748+
local_input_name = unique_name(inp)
749+
is_constant = isinstance(inp, Constant)
747750
input_storage = storage_map.setdefault(
748-
i, [None if not isinstance(i, Constant) else i.data]
751+
inp, [inp.data if is_constant else None]
749752
)
750-
if input_storage[0] is not None or isinstance(i, Constant):
753+
if (
754+
is_constant or input_storage[0] is not None
755+
) and inp not in tipifiyed_vars:
751756
# Constants need to be assigned locally and referenced
752757
global_env[local_input_name] = type_conversion_fn(
753-
input_storage[0], variable=i, storage=input_storage, **kwargs
758+
input_storage[0], variable=inp, storage=input_storage, **kwargs
754759
)
760+
tipifiyed_vars.add(inp)
755761
# TODO: We could attempt to use the storage arrays directly
762+
# Otherwise we're doubling the memory footprint of constants
756763
# E.g. `local_input_name = f"{local_input_name}[0]"`
757764
node_input_names.append(local_input_name)
758765

0 commit comments

Comments
 (0)