Make zips strict in pytensor/graph

Author: Armavica

pytensor/graph/basic.py

@@ -272,7 +272,7 @@ def clone_with_new_inputs(
         # as the output type depends on the input values and not just their types
         output_type_depends_on_input_value = self.op._output_type_depends_on_input_value
 
-        for i, (curr, new) in enumerate(zip(self.inputs, new_inputs)):
+        for i, (curr, new) in enumerate(zip(self.inputs, new_inputs, strict=True)):
             # Check if the input type changed or if the Op has output types that depend on input values
             if (curr.type != new.type) or output_type_depends_on_input_value:
                 # In strict mode, the cloned graph is assumed to be mathematically equivalent to the original one.
@@ -1295,7 +1295,7 @@ def clone_node_and_cache(
     if new_node.op is not node.op:
         clone_d.setdefault(node.op, new_node.op)
 
-    for old_o, new_o in zip(node.outputs, new_node.outputs):
+    for old_o, new_o in zip(node.outputs, new_node.outputs, strict=True):
         clone_d.setdefault(old_o, new_o)
 
     return new_node
@@ -1884,7 +1884,7 @@ def equal_computations(
     if in_ys is None:
         in_ys = []
 
-    for x, y in zip(xs, ys):
+    for x, y in zip(xs, ys, strict=True):
         if not isinstance(x, Variable) and not isinstance(y, Variable):
             return np.array_equal(x, y)
         if not isinstance(x, Variable):
@@ -1907,13 +1907,13 @@ def equal_computations(
     if len(in_xs) != len(in_ys):
         return False
 
-    for _x, _y in zip(in_xs, in_ys):
+    for _x, _y in zip(in_xs, in_ys, strict=True):
         if not (_y.type.in_same_class(_x.type)):
             return False
 
-    common = set(zip(in_xs, in_ys))
+    common = set(zip(in_xs, in_ys, strict=True))
     different: set[tuple[Variable, Variable]] = set()
-    for dx, dy in zip(xs, ys):
+    for dx, dy in zip(xs, ys, strict=True):
         assert isinstance(dx, Variable)
         # We checked above that both dx and dy have an owner or not
         if dx.owner is None:
@@ -1949,7 +1949,7 @@ def compare_nodes(nd_x, nd_y, common, different):
             return False
         else:
             all_in_common = True
-            for dx, dy in zip(nd_x.outputs, nd_y.outputs):
+            for dx, dy in zip(nd_x.outputs, nd_y.outputs, strict=True):
                 if (dx, dy) in different:
                     return False
                 if (dx, dy) not in common:
@@ -1959,7 +1959,7 @@ def compare_nodes(nd_x, nd_y, common, different):
                 return True
 
             # Compare the individual inputs for equality
-            for dx, dy in zip(nd_x.inputs, nd_y.inputs):
+            for dx, dy in zip(nd_x.inputs, nd_y.inputs, strict=True):
                 if (dx, dy) not in common:
                     # Equality between the variables is unknown, compare
                     # their respective owners, if they have some
@@ -1994,7 +1994,7 @@ def compare_nodes(nd_x, nd_y, common, different):
             # If the code reaches this statement then the inputs are pair-wise
             # equivalent so the outputs of the current nodes are also
             # pair-wise equivalents
-            for dx, dy in zip(nd_x.outputs, nd_y.outputs):
+            for dx, dy in zip(nd_x.outputs, nd_y.outputs, strict=True):
                 common.add((dx, dy))
 
             return True

pytensor/graph/op.py

@@ -231,14 +231,14 @@ def make_node(self, *inputs: Variable) -> Apply:
             )
         if not all(
             expected_type.is_super(var.type)
-            for var, expected_type in zip(inputs, self.itypes)
+            for var, expected_type in zip(inputs, self.itypes, strict=True)
         ):
             raise TypeError(
                 f"Invalid input types for Op {self}:\n"
                 + "\n".join(
                     f"Input {i}/{len(inputs)}: Expected {inp}, got {out}"
                     for i, (inp, out) in enumerate(
-                        zip(self.itypes, (inp.type for inp in inputs)),
+                        zip(self.itypes, (inp.type for inp in inputs), strict=True),
                         start=1,
                     )
                     if inp != out

pytensor/graph/replace.py

@@ -78,7 +78,7 @@ def clone_replace(
     items = list(_format_replace(replace).items())
 
     tmp_replace = [(x, x.type()) for x, y in items]
-    new_replace = [(x, y) for ((_, x), (_, y)) in zip(tmp_replace, items)]
+    new_replace = [(x, y) for ((_, x), (_, y)) in zip(tmp_replace, items, strict=True)]
     _, _outs, _ = rebuild_collect_shared(output, [], tmp_replace, [], **rebuild_kwds)
 
     # TODO Explain why we call it twice ?!
@@ -295,11 +295,11 @@ def vectorize_graph(
     inputs = truncated_graph_inputs(seq_outputs, ancestors_to_include=replace.keys())
     new_inputs = [replace.get(inp, inp) for inp in inputs]
 
-    vect_vars = dict(zip(inputs, new_inputs))
+    vect_vars = dict(zip(inputs, new_inputs, strict=True))
     for node in io_toposort(inputs, seq_outputs):
         vect_inputs = [vect_vars.get(inp, inp) for inp in node.inputs]
         vect_node = vectorize_node(node, *vect_inputs)
-        for output, vect_output in zip(node.outputs, vect_node.outputs):
+        for output, vect_output in zip(node.outputs, vect_node.outputs, strict=True):
             if output in vect_vars:
                 # This can happen when some outputs of a multi-output node are given a replacement,
                 # while some of the remaining outputs are still needed in the graph.

pytensor/graph/rewriting/basic.py

@@ -399,14 +399,14 @@ def print_profile(cls, stream, prof, level=0):
                 file=stream,
             )
         ll = []
-        for rewrite, nb_n in zip(rewrites, nb_nodes):
+        for rewrite, nb_n in zip(rewrites, nb_nodes, strict=True):
             if hasattr(rewrite, "__name__"):
                 name = rewrite.__name__
             else:
                 name = rewrite.name
             idx = rewrites.index(rewrite)
             ll.append((name, rewrite.__class__.__name__, idx, *nb_n))
-        lll = sorted(zip(prof, ll), key=lambda a: a[0])
+        lll = sorted(zip(prof, ll, strict=True), key=lambda a: a[0])
 
         for t, rewrite in lll[::-1]:
             i = rewrite[2]
@@ -480,7 +480,8 @@ def merge_profile(prof1, prof2):
 
         new_rewrite = SequentialGraphRewriter(*new_l)
         new_nb_nodes = [
-            (p1[0] + p2[0], p1[1] + p2[1]) for p1, p2 in zip(prof1[8], prof2[8])
+            (p1[0] + p2[0], p1[1] + p2[1])
+            for p1, p2 in zip(prof1[8], prof2[8], strict=True)
         ]
         new_nb_nodes.extend(prof1[8][len(new_nb_nodes) :])
         new_nb_nodes.extend(prof2[8][len(new_nb_nodes) :])
@@ -635,7 +636,7 @@ def process_node(self, fgraph, node):
 
             inputs_match = all(
                 node_in is cand_in
-                for node_in, cand_in in zip(node.inputs, candidate.inputs)
+                for node_in, cand_in in zip(node.inputs, candidate.inputs, strict=True)
             )
 
             if inputs_match and node.op == candidate.op:
@@ -649,6 +650,7 @@ def process_node(self, fgraph, node):
                         node.outputs,
                         candidate.outputs,
                         ["merge"] * len(node.outputs),
+                        strict=True,
                     )
                 )
 
@@ -721,7 +723,9 @@ def apply(self, fgraph):
                         inputs_match = all(
                             node_in is cand_in
                             for node_in, cand_in in zip(
-                                var.owner.inputs, candidate_var.owner.inputs
+                                var.owner.inputs,
+                                candidate_var.owner.inputs,
+                                strict=True,
                             )
                         )
 
@@ -1434,7 +1438,7 @@ def transform(self, fgraph, node):
         repl = self.op2.make_node(*node.inputs)
         if self.transfer_tags:
             repl.tag = copy.copy(node.tag)
-            for output, new_output in zip(node.outputs, repl.outputs):
+            for output, new_output in zip(node.outputs, repl.outputs, strict=True):
                 new_output.tag = copy.copy(output.tag)
         return repl.outputs
 
@@ -1616,7 +1620,7 @@ def transform(self, fgraph, node, get_nodes=True):
                     continue
                 ret = self.transform(fgraph, real_node, get_nodes=False)
                 if ret is not False and ret is not None:
-                    return dict(zip(real_node.outputs, ret))
+                    return dict(zip(real_node.outputs, ret, strict=True))
 
         if node.op != self.op:
             return False
@@ -1648,7 +1652,7 @@ def transform(self, fgraph, node, get_nodes=True):
                 len(node.outputs) == len(ret.owner.outputs)
                 and all(
                     o.type.is_super(new_o.type)
-                    for o, new_o in zip(node.outputs, ret.owner.outputs)
+                    for o, new_o in zip(node.outputs, ret.owner.outputs, strict=True)
                 )
             ):
                 return False
@@ -1940,7 +1944,7 @@ def process_node(
             )
         # None in the replacement mean that this variable isn't used
         # and we want to remove it
-        for r, rnew in zip(old_vars, replacements):
+        for r, rnew in zip(old_vars, replacements, strict=True):
             if rnew is None and len(fgraph.clients[r]) > 0:
                 raise ValueError(
                     f"Node rewriter {node_rewriter} tried to remove a variable"
@@ -1950,7 +1954,7 @@ def process_node(
         # the replacement
         repl_pairs = [
             (r, rnew)
-            for r, rnew in zip(old_vars, replacements)
+            for r, rnew in zip(old_vars, replacements, strict=True)
             if rnew is not r and rnew is not None
         ]
 
@@ -2633,17 +2637,23 @@ def print_profile(cls, stream, prof, level=0):
         print(blanc, "Global, final, and clean up rewriters", file=stream)
         for i in range(len(loop_timing)):
             print(blanc, f"Iter {int(i)}", file=stream)
-            for o, prof in zip(rewrite.global_rewriters, global_sub_profs[i]):
+            for o, prof in zip(
+                rewrite.global_rewriters, global_sub_profs[i], strict=True
+            ):
                 try:
                     o.print_profile(stream, prof, level + 2)
                 except NotImplementedError:
                     print(blanc, "merge not implemented for ", o)
-            for o, prof in zip(rewrite.final_rewriters, final_sub_profs[i]):
+            for o, prof in zip(
+                rewrite.final_rewriters, final_sub_profs[i], strict=True
+            ):
                 try:
                     o.print_profile(stream, prof, level + 2)
                 except NotImplementedError:
                     print(blanc, "merge not implemented for ", o)
-            for o, prof in zip(rewrite.cleanup_rewriters, cleanup_sub_profs[i]):
+            for o, prof in zip(
+                rewrite.cleanup_rewriters, cleanup_sub_profs[i], strict=True
+            ):
                 try:
                     o.print_profile(stream, prof, level + 2)
                 except NotImplementedError:
@@ -2861,7 +2871,7 @@ def local_recursive_function(
                     outs, rewritten_vars = local_recursive_function(
                         rewrite_list, inp, rewritten_vars, depth + 1
                     )
-                    for k, v in zip(inp.owner.outputs, outs):
+                    for k, v in zip(inp.owner.outputs, outs, strict=True):
                         rewritten_vars[k] = v
                     nw_in = outs[inp.owner.outputs.index(inp)]
 
@@ -2879,7 +2889,7 @@ def local_recursive_function(
             if ret is not False and ret is not None:
                 assert isinstance(ret, Sequence)
                 assert len(ret) == len(node.outputs), rewrite
-                for k, v in zip(node.outputs, ret):
+                for k, v in zip(node.outputs, ret, strict=True):
                     rewritten_vars[k] = v
                 results = ret
                 if ret[0].owner: