Make zips strict in pytensor/scalar

Author: Armavica

pytensor/scalar/basic.py

@@ -1147,7 +1147,7 @@ def perform(self, node, inputs, output_storage):
         else:
             variables = from_return_values(self.impl(*inputs))
             assert len(variables) == len(output_storage)
-            for storage, variable in zip(output_storage, variables):
+            for storage, variable in zip(output_storage, variables, strict=True):
                 storage[0] = variable
 
     def impl(self, *inputs):
@@ -4109,7 +4109,9 @@ def c_support_code(self, **kwargs):
 
     def c_support_code_apply(self, node, name):
         rval = []
-        for subnode, subnodename in zip(self.fgraph.toposort(), self.nodenames):
+        for subnode, subnodename in zip(
+            self.fgraph.toposort(), self.nodenames, strict=True
+        ):
             subnode_support_code = subnode.op.c_support_code_apply(
                 subnode, subnodename % dict(nodename=name)
             )
@@ -4215,7 +4217,7 @@ def __init__(self, inputs, outputs, name="Composite"):
             res2 = pytensor.compile.rebuild_collect_shared(
                 inputs=outputs[0].owner.op.inputs,
                 outputs=outputs[0].owner.op.outputs,
-                replace=dict(zip(outputs[0].owner.op.inputs, res[1])),
+                replace=dict(zip(outputs[0].owner.op.inputs, res[1], strict=True)),
             )
             assert len(res2[1]) == len(outputs)
             assert len(res[0]) == len(inputs)
@@ -4301,7 +4303,7 @@ def make_node(self, *inputs):
             assert len(inputs) == self.nin
             res = pytensor.compile.rebuild_collect_shared(
                 self.outputs,
-                replace=dict(zip(self.inputs, inputs)),
+                replace=dict(zip(self.inputs, inputs, strict=True)),
                 rebuild_strict=False,
             )
             # After rebuild_collect_shared, the Variable in inputs
@@ -4314,7 +4316,7 @@ def make_node(self, *inputs):
 
     def perform(self, node, inputs, output_storage):
         outputs = self.py_perform_fn(*inputs)
-        for storage, out_val in zip(output_storage, outputs):
+        for storage, out_val in zip(output_storage, outputs, strict=True):
             storage[0] = out_val
 
     def grad(self, inputs, output_grads):
@@ -4384,8 +4386,8 @@ def c_code_template(self):
     def c_code(self, node, nodename, inames, onames, sub):
         d = dict(
             chain(
-                zip((f"i{int(i)}" for i in range(len(inames))), inames),
-                zip((f"o{int(i)}" for i in range(len(onames))), onames),
+                zip((f"i{int(i)}" for i in range(len(inames))), inames, strict=True),
+                zip((f"o{int(i)}" for i in range(len(onames))), onames, strict=True),
             ),
             **sub,
         )
@@ -4433,7 +4435,7 @@ def apply(self, fgraph):
             )
             # make sure we don't produce any float16.
             assert not any(o.dtype == "float16" for o in new_node.outputs)
-            for o, no in zip(node.outputs, new_node.outputs):
+            for o, no in zip(node.outputs, new_node.outputs, strict=True):
                 mapping[o] = no
 
         new_ins = [mapping[inp] for inp in fgraph.inputs]
@@ -4477,7 +4479,7 @@ def handle_composite(node, mapping):
     new_op = node.op.clone_float32()
     new_outs = new_op(*[mapping[i] for i in node.inputs], return_list=True)
     assert len(new_outs) == len(node.outputs)
-    for o, no in zip(node.outputs, new_outs):
+    for o, no in zip(node.outputs, new_outs, strict=True):
         mapping[o] = no
 
 

pytensor/scalar/loop.py

@@ -93,7 +93,7 @@ def _validate_updates(
                 )
         else:
             update = outputs
-        for i, u in zip(init, update):
+        for i, u in zip(init, update, strict=False):
             if i.type != u.type:
                 raise TypeError(
                     "Init and update types must be the same: "
@@ -166,7 +166,7 @@ def make_node(self, n_steps, *inputs):
             # Make a new op with the right input types.
             res = rebuild_collect_shared(
                 self.outputs,
-                replace=dict(zip(self.inputs, inputs)),
+                replace=dict(zip(self.inputs, inputs, strict=True)),
                 rebuild_strict=False,
             )
             if self.is_while:
@@ -207,7 +207,7 @@ def perform(self, node, inputs, output_storage):
             for i in range(n_steps):
                 carry = inner_fn(*carry, *constant)
 
-        for storage, out_val in zip(output_storage, carry):
+        for storage, out_val in zip(output_storage, carry, strict=True):
             storage[0] = out_val
 
     @property
@@ -295,7 +295,7 @@ def c_code_template(self):
 
         # Set the carry variables to the output variables
         _c_code += "\n"
-        for init, update in zip(carry_subd.values(), update_subd.values()):
+        for init, update in zip(carry_subd.values(), update_subd.values(), strict=True):
             _c_code += f"{init} = {update};\n"
 
         # _c_code += 'printf("%%ld\\n", i);\n'
@@ -321,8 +321,8 @@ def c_code_template(self):
     def c_code(self, node, nodename, inames, onames, sub):
         d = dict(
             chain(
-                zip((f"i{int(i)}" for i in range(len(inames))), inames),
-                zip((f"o{int(i)}" for i in range(len(onames))), onames),
+                zip((f"i{int(i)}" for i in range(len(inames))), inames, strict=True),
+                zip((f"o{int(i)}" for i in range(len(onames))), onames, strict=True),
             ),
             **sub,
         )