Cleanup Function.__call__

Author: ricardoV94

pytensor/compile/function/types.py

@@ -326,8 +326,8 @@ class Function:
     def __init__(
         self,
         vm: "VM",
-        input_storage,
-        output_storage,
+        input_storage: list[Container],
+        output_storage: list[Container],
         indices,
         outputs,
         defaults,
@@ -388,6 +388,8 @@ def __init__(
         self.nodes_with_inner_function = []
         self.output_keys = output_keys
 
+        assert len(self.output_storage) == len(self.maker.fgraph.outputs)
+
         # See if we have any mutable / borrow inputs
         # TODO: this only need to be set if there is more than one input
         self._check_for_aliased_inputs = False
@@ -408,11 +410,6 @@ def __init__(
         finder = {}
         inv_finder = {}
 
-        def distribute(indices, cs, value):
-            input.distribute(value, indices, cs)
-            for c in cs:
-                c.provided += 1
-
         # Store the list of names of named inputs.
         named_inputs = []
         # Count the number of un-named inputs.
@@ -777,6 +774,13 @@ def checkSV(sv_ori, sv_rpl):
         f_cpy.maker.fgraph.name = name
         return f_cpy
 
+    def _restore_defaults(self):
+        for i, (required, refeed, value) in enumerate(self.defaults):
+            if refeed:
+                if isinstance(value, Container):
+                    value = value.storage[0]
+                self[i] = value
+
     def __call__(self, *args, **kwargs):
         """
         Evaluates value of a function on given arguments.
@@ -814,43 +818,42 @@ def restore_defaults():
                     self[i] = value
 
         profile = self.profile
-        t0 = time.perf_counter()
+        if profile:
+            t0 = time.perf_counter()
 
         output_subset = kwargs.pop("output_subset", None)
         if output_subset is not None and self.output_keys is not None:
             output_subset = [self.output_keys.index(key) for key in output_subset]
 
         # Reinitialize each container's 'provided' counter
         if self.trust_input:
-            i = 0
-            for arg in args:
-                s = self.input_storage[i]
-                s.storage[0] = arg
-                i += 1
+            for arg_container, arg in zip(self.input_storage, args, strict=False):
+                arg_container.storage[0] = arg
         else:
-            for c in self.input_storage:
-                c.provided = 0
+            for arg_container in self.input_storage:
+                arg_container.provided = 0
 
             if len(args) + len(kwargs) > len(self.input_storage):
                 raise TypeError("Too many parameter passed to pytensor function")
 
             # Set positional arguments
-            i = 0
-            for arg in args:
+            for arg_container, arg in zip(self.input_storage, args, strict=False):
                 # TODO: provide a option for skipping the filter if we really
                 # want speed.
-                s = self.input_storage[i]
                 # see this emails for a discuation about None as input
                 # https://groups.google.com/group/theano-dev/browse_thread/thread/920a5e904e8a8525/4f1b311a28fc27e5
                 if arg is None:
-                    s.storage[0] = arg
+                    arg_container.storage[0] = arg
                 else:
                     try:
-                        s.storage[0] = s.type.filter(
-                            arg, strict=s.strict, allow_downcast=s.allow_downcast
+                        arg_container.storage[0] = arg_container.type.filter(
+                            arg,
+                            strict=arg_container.strict,
+                            allow_downcast=arg_container.allow_downcast,
                         )
 
                     except Exception as e:
+                        i = self.input_storage.index(arg_container)
                         function_name = "pytensor function"
                         argument_name = "argument"
                         if self.name:
@@ -875,27 +878,23 @@ def restore_defaults():
                                 + function_name
                                 + f" at index {int(i)} (0-based). {where}"
                             ) + e.args
-                        restore_defaults()
+                        self._restore_defaults()
                         raise
-                s.provided += 1
-                i += 1
+                arg_container.provided += 1
 
         # Set keyword arguments
         if kwargs:  # for speed, skip the items for empty kwargs
             for k, arg in kwargs.items():
                 self[k] = arg
 
-        if (
-            not self.trust_input
-            and
-            # The getattr is only needed for old pickle
-            getattr(self, "_check_for_aliased_inputs", True)
-        ):
+        if not self.trust_input and self._check_for_aliased_inputs:
             # Collect aliased inputs among the storage space
             args_share_memory = []
-            for i in range(len(self.input_storage)):
-                i_var = self.maker.inputs[i].variable
-                i_val = self.input_storage[i].storage[0]
+            for i, (inp, arg_storage) in enumerate(
+                zip(self.maker.inputs, self.input_storage)
+            ):
+                i_var = inp.variable
+                i_val = arg_storage.storage[0]
                 if hasattr(i_var.type, "may_share_memory"):
                     is_aliased = False
                     for j in range(len(args_share_memory)):
@@ -932,36 +931,36 @@ def restore_defaults():
                             self.input_storage[j].storage[0]
                         )
 
-        # Check if inputs are missing, or if inputs were set more than once, or
-        # if we tried to provide inputs that are supposed to be implicit.
-        if not self.trust_input:
-            for c in self.input_storage:
-                if c.required and not c.provided:
-                    restore_defaults()
+            # Check if inputs are missing, or if inputs were set more than once, or
+            # if we tried to provide inputs that are supposed to be implicit.
+            for arg_container in self.input_storage:
+                if arg_container.required and not arg_container.provided:
+                    self._restore_defaults()
                     raise TypeError(
-                        f"Missing required input: {getattr(self.inv_finder[c], 'variable', self.inv_finder[c])}"
+                        f"Missing required input: {getattr(self.inv_finder[arg_container], 'variable', self.inv_finder[arg_container])}"
                     )
-                if c.provided > 1:
-                    restore_defaults()
+                if arg_container.provided > 1:
+                    self._restore_defaults()
                     raise TypeError(
-                        f"Multiple values for input: {getattr(self.inv_finder[c], 'variable', self.inv_finder[c])}"
+                        f"Multiple values for input: {getattr(self.inv_finder[arg_container], 'variable', self.inv_finder[arg_container])}"
                     )
-                if c.implicit and c.provided > 0:
-                    restore_defaults()
+                if arg_container.implicit and arg_container.provided > 0:
+                    self._restore_defaults()
                     raise TypeError(
-                        f"Tried to provide value for implicit input: {getattr(self.inv_finder[c], 'variable', self.inv_finder[c])}"
+                        f"Tried to provide value for implicit input: {getattr(self.inv_finder[arg_container], 'variable', self.inv_finder[arg_container])}"
                     )
 
         # Do the actual work
-        t0_fn = time.perf_counter()
+        if profile:
+            t0_fn = time.perf_counter()
         try:
             outputs = (
                 self.vm()
                 if output_subset is None
                 else self.vm(output_subset=output_subset)
             )
         except Exception:
-            restore_defaults()
+            self._restore_defaults()
             if hasattr(self.vm, "position_of_error"):
                 # this is a new vm-provided function or c linker
                 # they need this because the exception manipulation
@@ -979,9 +978,9 @@ def restore_defaults():
                 # old-style linkers raise their own exceptions
                 raise
 
-        dt_fn = time.perf_counter() - t0_fn
-        self.maker.mode.fn_time += dt_fn
         if profile:
+            dt_fn = time.perf_counter() - t0_fn
+            self.maker.mode.fn_time += dt_fn
             profile.vm_call_time += dt_fn
 
         # Retrieve the values that were computed
@@ -991,14 +990,13 @@ def restore_defaults():
 
         # Remove internal references to required inputs.
         # These cannot be re-used anyway.
-        for c in self.input_storage:
-            if c.required:
-                c.storage[0] = None
+        for arg_container in self.input_storage:
+            if arg_container.required:
+                arg_container.storage[0] = None
 
         # if we are allowing garbage collection, remove the
         # output reference from the internal storage cells
         if getattr(self.vm, "allow_gc", False):
-            assert len(self.output_storage) == len(self.maker.fgraph.outputs)
             for o_container, o_variable in zip(
                 self.output_storage, self.maker.fgraph.outputs
             ):
@@ -1020,17 +1018,12 @@ def restore_defaults():
             outputs = outputs[: self.n_returned_outputs]
 
         # Put default values back in the storage
-        restore_defaults()
-        #
-        # NOTE: This logic needs to be replicated in
-        #       scan.
-        #       grep for 'PROFILE_CODE'
-        #
-
-        dt_call = time.perf_counter() - t0
-        pytensor.compile.profiling.total_fct_exec_time += dt_call
-        self.maker.mode.call_time += dt_call
+        self._restore_defaults()
+
         if profile:
+            dt_call = time.perf_counter() - t0
+            pytensor.compile.profiling.total_fct_exec_time += dt_call
+            self.maker.mode.call_time += dt_call
             profile.fct_callcount += 1
             profile.fct_call_time += dt_call
             if hasattr(self.vm, "update_profile"):