Make the remaining zips strict

Author: Armavica

pytensor/gradient.py

@@ -213,7 +213,7 @@ def Rop(
 
     # Check that each element of wrt corresponds to an element
     # of eval_points with the same dimensionality.
-    for i, (wrt_elem, eval_point) in enumerate(zip(_wrt, _eval_points)):
+    for i, (wrt_elem, eval_point) in enumerate(zip(_wrt, _eval_points, strict=True)):
         try:
             if wrt_elem.type.ndim != eval_point.type.ndim:
                 raise ValueError(
@@ -262,7 +262,7 @@ def _traverse(node):
                     seen_nodes[inp.owner][inp.owner.outputs.index(inp)]
                 )
         same_type_eval_points = []
-        for x, y in zip(inputs, local_eval_points):
+        for x, y in zip(inputs, local_eval_points, strict=True):
             if y is not None:
                 if not isinstance(x, Variable):
                     x = pytensor.tensor.as_tensor_variable(x)
@@ -399,7 +399,7 @@ def Lop(
         _wrt = [pytensor.tensor.as_tensor_variable(x) for x in wrt]
 
     assert len(_f) == len(grads)
-    known = dict(zip(_f, grads))
+    known = dict(zip(_f, grads, strict=True))
 
     ret = grad(
         cost=None,
@@ -778,7 +778,7 @@ def subgraph_grad(wrt, end, start=None, cost=None, details=False):
             for i in range(len(grads)):
                 grads[i] += cost_grads[i]
 
-    pgrads = dict(zip(params, grads))
+    pgrads = dict(zip(params, grads, strict=True))
     # separate wrt from end grads:
     wrt_grads = [pgrads[k] for k in wrt]
     end_grads = [pgrads[k] for k in end]
@@ -1045,7 +1045,7 @@ def access_term_cache(node):
                     in [
                         input_to_output and output_to_cost
                         for input_to_output, output_to_cost in zip(
-                            input_to_outputs, outputs_connected
+                            input_to_outputs, outputs_connected, strict=True
                         )
                     ]
                 )
@@ -1071,7 +1071,7 @@ def access_term_cache(node):
                     not in [
                         in_to_out and out_to_cost and not out_nan
                         for in_to_out, out_to_cost, out_nan in zip(
-                            in_to_outs, outputs_connected, ograd_is_nan
+                            in_to_outs, outputs_connected, ograd_is_nan, strict=True
                         )
                     ]
                 )
@@ -1131,7 +1131,7 @@ def try_to_copy_if_needed(var):
                 # DO NOT force integer variables to have integer dtype.
                 # This is a violation of the op contract.
                 new_output_grads = []
-                for o, og in zip(node.outputs, output_grads):
+                for o, og in zip(node.outputs, output_grads, strict=True):
                     o_dt = getattr(o.type, "dtype", None)
                     og_dt = getattr(og.type, "dtype", None)
                     if (
@@ -1145,7 +1145,7 @@ def try_to_copy_if_needed(var):
 
                 # Make sure that, if new_output_grads[i] has a floating point
                 # dtype, it is the same dtype as outputs[i]
-                for o, ng in zip(node.outputs, new_output_grads):
+                for o, ng in zip(node.outputs, new_output_grads, strict=True):
                     o_dt = getattr(o.type, "dtype", None)
                     ng_dt = getattr(ng.type, "dtype", None)
                     if (
@@ -1167,7 +1167,9 @@ def try_to_copy_if_needed(var):
                 # by the user, not computed by Op.grad, and some gradients are
                 # only computed and returned, but never passed as another
                 # node's output grads.
-                for idx, packed in enumerate(zip(node.outputs, new_output_grads)):
+                for idx, packed in enumerate(
+                    zip(node.outputs, new_output_grads, strict=True)
+                ):
                     orig_output, new_output_grad = packed
                     if not hasattr(orig_output, "shape"):
                         continue
@@ -1233,7 +1235,7 @@ def try_to_copy_if_needed(var):
                     not in [
                         in_to_out and out_to_cost and not out_int
                         for in_to_out, out_to_cost, out_int in zip(
-                            in_to_outs, outputs_connected, output_is_int
+                            in_to_outs, outputs_connected, output_is_int, strict=True
                         )
                     ]
                 )
@@ -1314,7 +1316,7 @@ def try_to_copy_if_needed(var):
             # Check that op.connection_pattern matches the connectivity
             # logic driving the op.grad method
             for i, (ipt, ig, connected) in enumerate(
-                zip(inputs, input_grads, inputs_connected)
+                zip(inputs, input_grads, inputs_connected, strict=True)
             ):
                 actually_connected = not isinstance(ig.type, DisconnectedType)
 
@@ -1601,7 +1603,7 @@ def abs_rel_errors(self, g_pt):
         if len(g_pt) != len(self.gf):
             raise ValueError("argument has wrong number of elements", len(g_pt))
         errs = []
-        for i, (a, b) in enumerate(zip(g_pt, self.gf)):
+        for i, (a, b) in enumerate(zip(g_pt, self.gf, strict=True)):
             if a.shape != b.shape:
                 raise ValueError(
                     f"argument element {i} has wrong shapes {a.shape}, {b.shape}"

pytensor/ifelse.py

@@ -170,7 +170,9 @@ def make_node(self, condition: "TensorLike", *true_false_branches: Any):
         output_vars = []
         new_inputs_true_branch = []
         new_inputs_false_branch = []
-        for input_t, input_f in zip(inputs_true_branch, inputs_false_branch):
+        for input_t, input_f in zip(
+            inputs_true_branch, inputs_false_branch, strict=True
+        ):
             if not isinstance(input_t, Variable):
                 input_t = as_symbolic(input_t)
             if not isinstance(input_f, Variable):
@@ -207,7 +209,9 @@ def make_node(self, condition: "TensorLike", *true_false_branches: Any):
                 # allowed to have distinct shapes from either branch
                 new_shape = tuple(
                     s_t if s_t == s_f else None
-                    for s_t, s_f in zip(input_t.type.shape, input_f.type.shape)
+                    for s_t, s_f in zip(
+                        input_t.type.shape, input_f.type.shape, strict=True
+                    )
                 )
                 # TODO FIXME: The presence of this keyword is a strong
                 # assumption.  Find something that's guaranteed by the/a
@@ -301,7 +305,7 @@ def thunk():
                     if len(ls) > 0:
                         return ls
                     else:
-                        for out, t in zip(outputs, input_true_branch):
+                        for out, t in zip(outputs, input_true_branch, strict=True):
                             compute_map[out][0] = 1
                             val = storage_map[t][0]
                             if self.as_view:
@@ -321,7 +325,7 @@ def thunk():
                     if len(ls) > 0:
                         return ls
                     else:
-                        for out, f in zip(outputs, inputs_false_branch):
+                        for out, f in zip(outputs, inputs_false_branch, strict=True):
                             compute_map[out][0] = 1
                             # can't view both outputs unless destroyhandler
                             # improves
@@ -637,7 +641,7 @@ def apply(self, fgraph):
                     old_outs += [proposal.outputs]
                 else:
                     old_outs += proposal.outputs
-                pairs = list(zip(old_outs, new_outs))
+                pairs = list(zip(old_outs, new_outs, strict=True))
                 fgraph.replace_all_validate(pairs, reason="cond_merge")
 
 
@@ -736,7 +740,7 @@ def cond_merge_random_op(fgraph, main_node):
                 old_outs += [proposal.outputs]
             else:
                 old_outs += proposal.outputs
-            pairs = list(zip(old_outs, new_outs))
+            pairs = list(zip(old_outs, new_outs, strict=True))
             main_outs = clone_replace(main_node.outputs, replace=pairs)
             return main_outs
 

pytensor/printing.py

@@ -311,7 +311,7 @@ def debugprint(
         )
 
     for var, profile, storage_map, topo_order in zip(
-        outputs_to_print, profile_list, storage_maps, topo_orders
+        outputs_to_print, profile_list, storage_maps, topo_orders, strict=True
     ):
         if hasattr(var.owner, "op"):
             if (
@@ -936,7 +936,7 @@ def pp_process(input, new_precedence):
             str(i): x
             for i, x in enumerate(
                 pp_process(input, precedence)
-                for input, precedence in zip(node.inputs, precedences)
+                for input, precedence in zip(node.inputs, precedences, strict=True)
             )
         }
         r = pattern % d
@@ -1449,7 +1449,7 @@ def apply_name(node):
     if isinstance(fct, Function):
         # TODO: Get rid of all this `expanded_inputs` nonsense and use
         # `fgraph.update_mapping`
-        function_inputs = zip(fct.maker.expanded_inputs, fgraph.inputs)
+        function_inputs = zip(fct.maker.expanded_inputs, fgraph.inputs, strict=True)
         for i, fg_ii in reversed(list(function_inputs)):
             if i.update is not None:
                 k = outputs.pop()