Make zips strict in pytensor/link

Author: Armavica

pytensor/link/basic.py

@@ -385,11 +385,11 @@ def make_all(
             f,
             [
                 Container(input, storage)
-                for input, storage in zip(fgraph.inputs, input_storage)
+                for input, storage in zip(fgraph.inputs, input_storage, strict=True)
             ],
             [
                 Container(output, storage, readonly=True)
-                for output, storage in zip(fgraph.outputs, output_storage)
+                for output, storage in zip(fgraph.outputs, output_storage, strict=True)
             ],
             thunks,
             order,
@@ -509,7 +509,9 @@ def make_thunk(self, **kwargs):
         kwargs.pop("input_storage", None)
         make_all += [x.make_all(**kwargs) for x in self.linkers[1:]]
 
-        fns, input_lists, output_lists, thunk_lists, order_lists = zip(*make_all)
+        fns, input_lists, output_lists, thunk_lists, order_lists = zip(
+            *make_all, strict=True
+        )
 
         order_list0 = order_lists[0]
         for order_list in order_lists[1:]:
@@ -521,12 +523,12 @@ def make_thunk(self, **kwargs):
         inputs0 = input_lists[0]
         outputs0 = output_lists[0]
 
-        thunk_groups = list(zip(*thunk_lists))
-        order = [x[0] for x in zip(*order_lists)]
+        thunk_groups = list(zip(*thunk_lists, strict=True))
+        order = [x[0] for x in zip(*order_lists, strict=True)]
 
         to_reset = [
             thunk.outputs[j]
-            for thunks, node in zip(thunk_groups, order)
+            for thunks, node in zip(thunk_groups, order, strict=True)
             for j, output in enumerate(node.outputs)
             if output in no_recycling
             for thunk in thunks
@@ -537,12 +539,12 @@ def make_thunk(self, **kwargs):
 
         def f():
             for inputs in input_lists[1:]:
-                for input1, input2 in zip(inputs0, inputs):
+                for input1, input2 in zip(inputs0, inputs, strict=True):
                     input2.storage[0] = copy(input1.storage[0])
             for x in to_reset:
                 x[0] = None
             pre(self, [input.data for input in input_lists[0]], order, thunk_groups)
-            for i, (thunks, node) in enumerate(zip(thunk_groups, order)):
+            for i, (thunks, node) in enumerate(zip(thunk_groups, order, strict=True)):
                 try:
                     wrapper(self.fgraph, i, node, *thunks)
                 except Exception:
@@ -664,7 +666,9 @@ def thunk(
         ):
             outputs = fgraph_jit(*[self.input_filter(x[0]) for x in thunk_inputs])
 
-            for o_var, o_storage, o_val in zip(fgraph.outputs, thunk_outputs, outputs):
+            for o_var, o_storage, o_val in zip(
+                fgraph.outputs, thunk_outputs, outputs, strict=True
+            ):
                 compute_map[o_var][0] = True
                 o_storage[0] = self.output_filter(o_var, o_val)
             return outputs
@@ -730,11 +734,11 @@ def make_all(self, input_storage=None, output_storage=None, storage_map=None):
             fn,
             [
                 Container(input, storage)
-                for input, storage in zip(fgraph.inputs, input_storage)
+                for input, storage in zip(fgraph.inputs, input_storage, strict=True)
             ],
             [
                 Container(output, storage, readonly=True)
-                for output, storage in zip(fgraph.outputs, output_storage)
+                for output, storage in zip(fgraph.outputs, output_storage, strict=True)
             ],
             thunks,
             nodes,

pytensor/link/utils.py

@@ -88,7 +88,7 @@ def map_storage(
         assert len(fgraph.inputs) == len(input_storage)
 
     # add input storage into storage_map
-    for r, storage in zip(fgraph.inputs, input_storage):
+    for r, storage in zip(fgraph.inputs, input_storage, strict=True):
         if r in storage_map:
             assert storage_map[r] is storage, (
                 "Given input_storage conflicts "
@@ -108,7 +108,7 @@ def map_storage(
     # allocate output storage
     if output_storage is not None:
         assert len(fgraph.outputs) == len(output_storage)
-        for r, storage in zip(fgraph.outputs, output_storage):
+        for r, storage in zip(fgraph.outputs, output_storage, strict=True):
             if r in storage_map:
                 assert storage_map[r] is storage, (
                     "Given output_storage confl"
@@ -191,7 +191,7 @@ def streamline_default_f():
                 x[0] = None
             try:
                 for thunk, node, old_storage in zip(
-                    thunks, order, post_thunk_old_storage
+                    thunks, order, post_thunk_old_storage, strict=True
                 ):
                     thunk()
                     for old_s in old_storage:
@@ -206,7 +206,7 @@ def streamline_nice_errors_f():
             for x in no_recycling:
                 x[0] = None
             try:
-                for thunk, node in zip(thunks, order):
+                for thunk, node in zip(thunks, order, strict=True):
                     thunk()
             except Exception:
                 raise_with_op(fgraph, node, thunk)

pytensor/link/vm.py

@@ -244,7 +244,7 @@ def clear_storage(self):
     def update_profile(self, profile):
         """Update a profile object."""
         for node, thunk, t, c in zip(
-            self.nodes, self.thunks, self.call_times, self.call_counts
+            self.nodes, self.thunks, self.call_times, self.call_counts, strict=True
         ):
             profile.apply_time[(self.fgraph, node)] += t
 
@@ -310,7 +310,9 @@ def __init__(
         self.output_storage = output_storage
         self.inp_storage_and_out_idx = tuple(
             (inp_storage, self.fgraph.outputs.index(update_vars[inp]))
-            for inp, inp_storage in zip(self.fgraph.inputs, self.input_storage)
+            for inp, inp_storage in zip(
+                self.fgraph.inputs, self.input_storage, strict=True
+            )
             if inp in update_vars
         )
 
@@ -1247,7 +1249,7 @@ def make_all(
             self.profile.linker_node_make_thunks += t1 - t0
             self.profile.linker_make_thunk_time = linker_make_thunk_time
 
-        for node, thunk in zip(order, thunks):
+        for node, thunk in zip(order, thunks, strict=True):
             thunk.inputs = [storage_map[v] for v in node.inputs]
             thunk.outputs = [storage_map[v] for v in node.outputs]
 
@@ -1303,11 +1305,11 @@ def make_all(
             vm,
             [
                 Container(input, storage)
-                for input, storage in zip(fgraph.inputs, input_storage)
+                for input, storage in zip(fgraph.inputs, input_storage, strict=True)
             ],
             [
                 Container(output, storage, readonly=True)
-                for output, storage in zip(fgraph.outputs, output_storage)
+                for output, storage in zip(fgraph.outputs, output_storage, strict=True)
             ],
             thunks,
             order,