Skip to content

Commit 0d9fa06

Browse files
committed
Avoid calling type_conversion_fn repeatedly on the same variables
1 parent 645adc9 commit 0d9fa06

File tree

1 file changed

+19
-5
lines changed

1 file changed

+19
-5
lines changed

pytensor/link/utils.py

Lines changed: 19 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,29 @@ 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,
752+
[
753+
inp.data # type: ignore[attr-defined]
754+
if is_constant
755+
else None
756+
],
749757
)
750-
if input_storage[0] is not None or isinstance(i, Constant):
758+
if (
759+
is_constant or input_storage[0] is not None
760+
) and inp not in tipifiyed_vars:
751761
# Constants need to be assigned locally and referenced
762+
# FIXME: This is converting shared variables, but these may change later,
763+
# so this one-time conversion is wasteful / not robust
752764
global_env[local_input_name] = type_conversion_fn(
753-
input_storage[0], variable=i, storage=input_storage, **kwargs
765+
input_storage[0], variable=inp, storage=input_storage, **kwargs
754766
)
767+
tipifiyed_vars.add(inp)
755768
# TODO: We could attempt to use the storage arrays directly
769+
# Otherwise we're doubling the memory footprint of constants
756770
# E.g. `local_input_name = f"{local_input_name}[0]"`
757771
node_input_names.append(local_input_name)
758772

0 commit comments

Comments
 (0)