Make zips strict in pytensor/link/c

Author: Armavica

pytensor/link/c/basic.py

@@ -1119,11 +1119,15 @@ def __compile__(
             module,
             [
                 Container(input, storage)
-                for input, storage in zip(self.fgraph.inputs, input_storage)
+                for input, storage in zip(
+                    self.fgraph.inputs, input_storage, strict=True
+                )
             ],
             [
                 Container(output, storage, readonly=True)
-                for output, storage in zip(self.fgraph.outputs, output_storage)
+                for output, storage in zip(
+                    self.fgraph.outputs, output_storage, strict=True
+                )
             ],
             error_storage,
         )
@@ -1891,11 +1895,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,
@@ -1993,22 +1997,26 @@ def make_thunk(self, **kwargs):
         )
 
         def f():
-            for input1, input2 in zip(i1, i2):
+            for input1, input2 in zip(i1, i2, strict=True):
                 # Set the inputs to be the same in both branches.
                 # The copy is necessary in order for inplace ops not to
                 # interfere.
                 input2.storage[0] = copy(input1.storage[0])
-            for thunk1, thunk2, node1, node2 in zip(thunks1, thunks2, order1, order2):
-                for output, storage in zip(node1.outputs, thunk1.outputs):
+            for thunk1, thunk2, node1, node2 in zip(
+                thunks1, thunks2, order1, order2, strict=True
+            ):
+                for output, storage in zip(node1.outputs, thunk1.outputs, strict=True):
                     if output in no_recycling:
                         storage[0] = None
-                for output, storage in zip(node2.outputs, thunk2.outputs):
+                for output, storage in zip(node2.outputs, thunk2.outputs, strict=True):
                     if output in no_recycling:
                         storage[0] = None
                 try:
                     thunk1()
                     thunk2()
-                    for output1, output2 in zip(thunk1.outputs, thunk2.outputs):
+                    for output1, output2 in zip(
+                        thunk1.outputs, thunk2.outputs, strict=True
+                    ):
                         self.checker(output1, output2)
                 except Exception:
                     raise_with_op(fgraph, node1)

pytensor/link/c/cmodule.py

@@ -2443,7 +2443,7 @@ def patch_ldflags(flag_list: list[str]) -> list[str]:
         if not libs:
             return flag_list
         libs = GCC_compiler.linking_patch(lib_dirs, libs)
-        for flag_idx, lib in zip(flag_idxs, libs):
+        for flag_idx, lib in zip(flag_idxs, libs, strict=True):
             flag_list[flag_idx] = lib
         return flag_list
 

pytensor/link/c/op.py

@@ -59,7 +59,7 @@ def make_c_thunk(
         e = FunctionGraph(node.inputs, node.outputs)
         e_no_recycling = [
             new_o
-            for (new_o, old_o) in zip(e.outputs, node.outputs)
+            for (new_o, old_o) in zip(e.outputs, node.outputs, strict=True)
             if old_o in no_recycling
         ]
         cl = pytensor.link.c.basic.CLinker().accept(e, no_recycling=e_no_recycling)

pytensor/link/c/params_type.py

@@ -709,7 +709,7 @@ def c_support_code(self, **kwargs):
         c_init_list = []
         c_cleanup_list = []
         c_extract_list = []
-        for attribute_name, type_instance in zip(self.fields, self.types):
+        for attribute_name, type_instance in zip(self.fields, self.types, strict=True):
             try:
                 # c_support_code() may return a code string or a list of code strings.
                 support_code = type_instance.c_support_code()