Make zips strict in pytensor/tensor/rewriting

Author: Armavica

pytensor/tensor/rewriting/basic.py

@@ -97,11 +97,11 @@ def broadcasted_by(x: TensorVariable, y: TensorVariable) -> bool:
     if len(bx) < len(by):
         return True
     bx = bx[-len(by) :]
-    return any(bx_dim and not by_dim for bx_dim, by_dim in zip(bx, by))
+    return any(bx_dim and not by_dim for bx_dim, by_dim in zip(bx, by, strict=True))
 
 
 def merge_broadcastables(broadcastables):
-    return [all(bcast) for bcast in zip(*broadcastables)]
+    return [all(bcast) for bcast in zip(*broadcastables, strict=True)]
 
 
 def alloc_like(
@@ -1203,7 +1203,7 @@ def local_merge_alloc(fgraph, node):
     # broadcasted dimensions to its inputs[0]. Eg:
     # Alloc(Alloc(m, y, 1, 1), x, y, z, w) -> Alloc(m, x, y, z, w)
     i = 0
-    for dim_inner, dim_outer in zip(dims_inner_rev, dims_outer_rev):
+    for dim_inner, dim_outer in zip(dims_inner_rev, dims_outer_rev, strict=True):
         if dim_inner != dim_outer:
             if isinstance(dim_inner, Constant) and dim_inner.data == 1:
                 pass

pytensor/tensor/rewriting/blas.py

@@ -502,7 +502,7 @@ def on_import(new_node):
                     ].tag.values_eq_approx = values_eq_approx_remove_inf_nan
                     try:
                         fgraph.replace_all_validate_remove(
-                            list(zip(node.outputs, new_outputs)),
+                            list(zip(node.outputs, new_outputs, strict=True)),
                             [old_dot22],
                             reason="GemmOptimizer",
                             # For now we disable the warning as we know case

pytensor/tensor/rewriting/blockwise.py

@@ -109,7 +109,7 @@ def local_blockwise_alloc(fgraph, node):
     new_inputs = []
     batch_shapes = []
     can_push_any_alloc = False
-    for inp, inp_sig in zip(node.inputs, op.inputs_sig):
+    for inp, inp_sig in zip(node.inputs, op.inputs_sig, strict=True):
         if inp.owner and isinstance(inp.owner.op, Alloc):
             # Push batch dims from Alloc
             value, *shape = inp.owner.inputs
@@ -130,6 +130,7 @@ def local_blockwise_alloc(fgraph, node):
                     for broadcastable, dim in zip(
                         squeezed_value.type.broadcastable[:squeezed_value_batch_ndim],
                         tuple(squeezed_value.shape)[:squeezed_value_batch_ndim],
+                        strict=True,
                     )
                 ]
                 squeezed_value = alloc(squeezed_value, *batch_shape, *core_shape)
@@ -143,7 +144,7 @@ def local_blockwise_alloc(fgraph, node):
                 tuple(
                     1 if broadcastable else dim
                     for broadcastable, dim in zip(
-                        inp.type.broadcastable, shape[:batch_ndim]
+                        inp.type.broadcastable, shape[:batch_ndim], strict=True
                     )
                 )
             )
@@ -166,7 +167,9 @@ def local_blockwise_alloc(fgraph, node):
         # We pick the most parsimonious batch dim from the pushed Alloc
         missing_ndim = old_out_type.ndim - new_out_type.ndim
         batch_shape = ([1] * missing_ndim + list(new_outs[0].shape))[:batch_ndim]
-        for i, batch_dims in enumerate(zip(*batch_shapes)):  # Transpose shape tuples
+        for i, batch_dims in enumerate(
+            zip(*batch_shapes, strict=True)
+        ):  # Transpose shape tuples
             for batch_dim in batch_dims:
                 if batch_dim == 1:
                     continue

pytensor/tensor/rewriting/elemwise.py

@@ -300,7 +300,7 @@ def apply(self, fgraph):
                         )
                         new_node = new_outputs[0].owner
 
-                        for r, new_r in zip(node.outputs, new_outputs):
+                        for r, new_r in zip(node.outputs, new_outputs, strict=True):
                             prof["nb_call_replace"] += 1
                             fgraph.replace(
                                 r, new_r, reason="inplace_elemwise_optimizer"
@@ -1037,12 +1037,12 @@ def update_fuseable_mappings_after_fg_replace(
             )
             if not isinstance(composite_outputs, list):
                 composite_outputs = [composite_outputs]
-            for old_out, composite_out in zip(outputs, composite_outputs):
+            for old_out, composite_out in zip(outputs, composite_outputs, strict=True):
                 if old_out.name:
                     composite_out.name = old_out.name
 
             fgraph.replace_all_validate(
-                list(zip(outputs, composite_outputs)),
+                list(zip(outputs, composite_outputs, strict=True)),
                 reason=self.__class__.__name__,
             )
             nb_replacement += 1
@@ -1118,7 +1118,7 @@ def local_useless_composite_outputs(fgraph, node):
         used_inputs = [node.inputs[i] for i in used_inputs_idxs]
         c = ps.Composite(inputs=used_inner_inputs, outputs=used_inner_outputs)
         e = Elemwise(scalar_op=c)(*used_inputs, return_list=True)
-        return dict(zip([node.outputs[i] for i in used_outputs_idxs], e))
+        return dict(zip([node.outputs[i] for i in used_outputs_idxs], e, strict=True))
 
 
 @node_rewriter([CAReduce])
@@ -1218,7 +1218,9 @@ def local_inline_composite_constants(fgraph, node):
     new_outer_inputs = []
     new_inner_inputs = []
     inner_replacements = {}
-    for outer_inp, inner_inp in zip(node.inputs, composite_op.fgraph.inputs):
+    for outer_inp, inner_inp in zip(
+        node.inputs, composite_op.fgraph.inputs, strict=True
+    ):
         # Complex variables don't have a `c_literal` that can be inlined
         if "complex" not in outer_inp.type.dtype:
             unique_value = get_unique_constant_value(outer_inp)
@@ -1355,7 +1357,7 @@ def local_useless_2f1grad_loop(fgraph, node):
 
     replacements = {converges: new_converges}
     i = 0
-    for grad_var, is_used in zip(grad_vars, grad_var_is_used):
+    for grad_var, is_used in zip(grad_vars, grad_var_is_used, strict=True):
         if not is_used:
             continue
         replacements[grad_var] = new_outs[i]

pytensor/tensor/rewriting/math.py

@@ -1142,7 +1142,9 @@ def transform(self, fgraph, node):
         num, denum = self.simplify(list(orig_num), list(orig_denum), out.type)
 
         def same(x, y):
-            return len(x) == len(y) and all(np.all(xe == ye) for xe, ye in zip(x, y))
+            return len(x) == len(y) and all(
+                np.all(xe == ye) for xe, ye in zip(x, y, strict=True)
+            )
 
         if (
             same(orig_num, num)
@@ -2445,7 +2447,9 @@ def distribute_greedy(pos_pairs, neg_pairs, num, denum, out_type, minscore=0):
             [(n + num, d + denum, out_type) for (n, d) in neg_pairs],
         )
     )
-    for (n, d), (nn, dd) in zip(pos_pairs + neg_pairs, new_pos_pairs + new_neg_pairs):
+    for (n, d), (nn, dd) in zip(
+        pos_pairs + neg_pairs, new_pos_pairs + new_neg_pairs, strict=True
+    ):
         # We calculate how many operations we are saving with the new
         # num and denum
         score += len(n) + div_cost * len(d) - len(nn) - div_cost * len(dd)

pytensor/tensor/rewriting/shape.py

@@ -186,7 +186,7 @@ def get_shape(self, var, idx):
 
             # Only change the variables and dimensions that would introduce
             # extra computation
-            for new_shps, out in zip(o_shapes, node.outputs):
+            for new_shps, out in zip(o_shapes, node.outputs, strict=True):
                 if not hasattr(out.type, "ndim"):
                     continue
 
@@ -578,7 +578,7 @@ def on_import(self, fgraph, node, reason):
                 new_shape += sh[len(new_shape) :]
                 o_shapes[sh_idx] = tuple(new_shape)
 
-        for r, s in zip(node.outputs, o_shapes):
+        for r, s in zip(node.outputs, o_shapes, strict=True):
             self.set_shape(r, s)
 
     def on_change_input(self, fgraph, node, i, r, new_r, reason):
@@ -709,7 +709,7 @@ def same_shape(
         sx = canon_shapes[: len(sx)]
         sy = canon_shapes[len(sx) :]
 
-        for dx, dy in zip(sx, sy):
+        for dx, dy in zip(sx, sy, strict=True):
             if not equal_computations([dx], [dy]):
                 return False
 
@@ -778,7 +778,7 @@ def f(fgraph, node):
         # rewrite.
         if rval.type.ndim == node.outputs[0].type.ndim and all(
             s1 == s1
-            for s1, s2 in zip(rval.type.shape, node.outputs[0].type.shape)
+            for s1, s2 in zip(rval.type.shape, node.outputs[0].type.shape, strict=True)
             if s1 == 1 or s2 == 1
         ):
             return [rval]
@@ -817,7 +817,7 @@ def local_useless_reshape(fgraph, node):
         and output.type.ndim == 1
         and all(
             s1 == s2
-            for s1, s2 in zip(inp.type.shape, output.type.shape)
+            for s1, s2 in zip(inp.type.shape, output.type.shape, strict=True)
             if s1 == 1 or s2 == 1
         )
     ):
@@ -1068,7 +1068,9 @@ def local_specify_shape_lift(fgraph, node):
 
             nonbcast_dims = {
                 i
-                for i, (dim, bcast) in enumerate(zip(shape, out_broadcastable))
+                for i, (dim, bcast) in enumerate(
+                    zip(shape, out_broadcastable, strict=True)
+                )
                 if (not bcast and not NoneConst.equals(dim))
             }
             new_elem_inps = elem_inps.copy()
@@ -1170,7 +1172,7 @@ def local_useless_dimshuffle_in_reshape(fgraph, node):
     new_order = node.inputs[0].owner.op.new_order
     inp = node.inputs[0].owner.inputs[0]
     new_order_of_nonbroadcast = []
-    for i, s in zip(new_order, node.inputs[0].type.shape):
+    for i, s in zip(new_order, node.inputs[0].type.shape, strict=True):
         if s != 1:
             new_order_of_nonbroadcast.append(i)
     no_change_in_order = all(
@@ -1197,7 +1199,7 @@ def local_useless_unbroadcast(fgraph, node):
         x = node.inputs[0]
         if x.type.ndim == node.outputs[0].type.ndim and all(
             s1 == s2
-            for s1, s2 in zip(x.type.shape, node.outputs[0].type.shape)
+            for s1, s2 in zip(x.type.shape, node.outputs[0].type.shape, strict=True)
             if s1 == 1 or s2 == 1
         ):
             # No broadcastable flag was modified

pytensor/tensor/rewriting/subtensor.py

@@ -647,7 +647,7 @@ def local_subtensor_of_alloc(fgraph, node):
     # Slices to take from val
     val_slices = []
 
-    for i, (sl, dim) in enumerate(zip(slices, dims)):
+    for i, (sl, dim) in enumerate(zip(slices, dims, strict=True)):
         # If val was not copied over that dim,
         # we need to take the appropriate subtensor on it.
         if i >= n_added_dims:
@@ -1771,7 +1771,7 @@ def local_join_subtensors(fgraph, node):
         if all(
             idxs_nonaxis_subtensor1 == idxs_nonaxis_subtensor2
             for i, (idxs_nonaxis_subtensor1, idxs_nonaxis_subtensor2) in enumerate(
-                zip(idxs_subtensor1, idxs_subtensor2)
+                zip(idxs_subtensor1, idxs_subtensor2, strict=True)
             )
             if i != axis
         ):
@@ -1913,7 +1913,7 @@ def local_blockwise_advanced_inc_subtensor(fgraph, node):
 
     x_batch_bcast = x.type.broadcastable[:batch_ndim]
     y_batch_bcast = y.type.broadcastable[:batch_ndim]
-    if any(xb and not yb for xb, yb in zip(x_batch_bcast, y_batch_bcast)):
+    if any(xb and not yb for xb, yb in zip(x_batch_bcast, y_batch_bcast, strict=True)):
         # Need to broadcast batch x dims
         batch_shape = tuple(
             x_dim if (not xb or yb) else y_dim
@@ -1922,6 +1922,7 @@ def local_blockwise_advanced_inc_subtensor(fgraph, node):
                 tuple(x.shape)[:batch_ndim],
                 y_batch_bcast,
                 tuple(y.shape)[:batch_ndim],
+                strict=True,
             )
         )
         core_shape = tuple(x.shape)[batch_ndim:]