Make zips strict in pytensor/tensor

Author: Armavica

pytensor/tensor/basic.py

@@ -1524,6 +1524,7 @@ def make_node(self, value, *shape):
                 extended_value_broadcastable,
                 extended_value_static_shape,
                 static_shape,
+                strict=True,
             )
         ):
             # If value is not broadcastable and we don't know the target static shape: use value static shape
@@ -1544,7 +1545,7 @@ def make_node(self, value, *shape):
     def _check_runtime_broadcast(node, value, shape):
         value_static_shape = node.inputs[0].type.shape
         for v_static_dim, value_dim, out_dim in zip(
-            value_static_shape[::-1], value.shape[::-1], shape[::-1]
+            value_static_shape[::-1], value.shape[::-1], shape[::-1], strict=False
         ):
             if v_static_dim is None and value_dim == 1 and out_dim != 1:
                 raise ValueError(Alloc._runtime_broadcast_error_msg)
@@ -1647,6 +1648,7 @@ def grad(self, inputs, grads):
                 inputs[0].type.shape,
                 # We need the dimensions corresponding to x
                 grads[0].type.shape[-inputs[0].ndim :],
+                strict=False,
             )
         ):
             if ib == 1 and gb != 1:
@@ -2157,7 +2159,7 @@ def grad(self, inputs, g_outputs):
             ]
         # Else, we have to make them zeros before joining them
         new_g_outputs = []
-        for o, g in zip(outputs, g_outputs):
+        for o, g in zip(outputs, g_outputs, strict=True):
             if isinstance(g.type, DisconnectedType):
                 new_g_outputs.append(o.zeros_like())
             else:
@@ -2586,7 +2588,7 @@ def grad(self, axis_and_tensors, grads):
                 else specify_broadcastable(
                     g, *(ax for (ax, s) in enumerate(t.type.shape) if s == 1)
                 )
-                for t, g in zip(tens, split_gz)
+                for t, g in zip(tens, split_gz, strict=True)
             ]
             rval = rval + split_gz
         else:
@@ -2694,7 +2696,7 @@ def vectorize_join(op: Join, node, batch_axis, *batch_inputs):
     ):
         batch_ndims = {
             batch_input.type.ndim - old_input.type.ndim
-            for batch_input, old_input in zip(batch_inputs, old_inputs)
+            for batch_input, old_input in zip(batch_inputs, old_inputs, strict=True)
         }
         if len(batch_ndims) == 1:
             [batch_ndim] = batch_ndims
@@ -3276,7 +3278,7 @@ def __getitem__(self, *args):
             tuple([1] * j + [r.shape[0]] + [1] * (ndim - 1 - j))
             for j, r in enumerate(ranges)
         ]
-        ranges = [r.reshape(shape) for r, shape in zip(ranges, shapes)]
+        ranges = [r.reshape(shape) for r, shape in zip(ranges, shapes, strict=True)]
         if self.sparse:
             grids = ranges
         else:
@@ -3348,7 +3350,7 @@ def make_node(self, x, y, inverse):
 
         out_shape = [
             1 if xb == 1 and yb == 1 else None
-            for xb, yb in zip(x.type.shape, y.type.shape)
+            for xb, yb in zip(x.type.shape, y.type.shape, strict=True)
         ]
         out_type = tensor(dtype=x.type.dtype, shape=out_shape)
 
@@ -3413,7 +3415,7 @@ def perform(self, node, inp, out):
 
         # Make sure the output is big enough
         out_s = []
-        for xdim, ydim in zip(x_s, y_s):
+        for xdim, ydim in zip(x_s, y_s, strict=True):
             if xdim == ydim:
                 outdim = xdim
             elif xdim == 1:
@@ -3473,7 +3475,7 @@ def grad(self, inp, grads):
         assert gx.type.ndim == x.type.ndim
         assert all(
             s1 == s2
-            for s1, s2 in zip(gx.type.shape, x.type.shape)
+            for s1, s2 in zip(gx.type.shape, x.type.shape, strict=True)
             if s1 == 1 or s2 == 1
         )
 
@@ -3983,7 +3985,7 @@ def moveaxis(
 
     order = [n for n in range(a.ndim) if n not in source]
 
-    for dest, src in sorted(zip(destination, source)):
+    for dest, src in sorted(zip(destination, source, strict=True)):
         order.insert(dest, src)
 
     result = a.dimshuffle(order)
@@ -4337,7 +4339,7 @@ def _make_along_axis_idx(arr_shape, indices, axis):
     # build a fancy index, consisting of orthogonal aranges, with the
     # requested index inserted at the right location
     fancy_index = []
-    for dim, n in zip(dest_dims, arr_shape):
+    for dim, n in zip(dest_dims, arr_shape, strict=True):
         if dim is None:
             fancy_index.append(indices)
         else:

pytensor/tensor/blockwise.py

@@ -88,7 +88,7 @@ def __getstate__(self):
 
     def _create_dummy_core_node(self, inputs: Sequence[TensorVariable]) -> Apply:
         core_input_types = []
-        for i, (inp, sig) in enumerate(zip(inputs, self.inputs_sig)):
+        for i, (inp, sig) in enumerate(zip(inputs, self.inputs_sig, strict=True)):
             if inp.type.ndim < len(sig):
                 raise ValueError(
                     f"Input {i} {inp} has insufficient core dimensions for signature {self.signature}"
@@ -106,7 +106,9 @@ def _create_dummy_core_node(self, inputs: Sequence[TensorVariable]) -> Apply:
             raise ValueError(
                 f"Insufficient number of outputs for signature {self.signature}: {len(core_node.outputs)}"
             )
-        for i, (core_out, sig) in enumerate(zip(core_node.outputs, self.outputs_sig)):
+        for i, (core_out, sig) in enumerate(
+            zip(core_node.outputs, self.outputs_sig, strict=True)
+        ):
             if core_out.type.ndim != len(sig):
                 raise ValueError(
                     f"Output {i} of {self.core_op} has wrong number of core dimensions for signature {self.signature}: {core_out.type.ndim}"
@@ -120,12 +122,13 @@ def make_node(self, *inputs):
         core_node = self._create_dummy_core_node(inputs)
 
         batch_ndims = max(
-            inp.type.ndim - len(sig) for inp, sig in zip(inputs, self.inputs_sig)
+            inp.type.ndim - len(sig)
+            for inp, sig in zip(inputs, self.inputs_sig, strict=True)
         )
 
         batched_inputs = []
         batch_shapes = []
-        for i, (inp, sig) in enumerate(zip(inputs, self.inputs_sig)):
+        for i, (inp, sig) in enumerate(zip(inputs, self.inputs_sig, strict=True)):
             # Append missing dims to the left
             missing_batch_ndims = batch_ndims - (inp.type.ndim - len(sig))
             if missing_batch_ndims:
@@ -141,7 +144,7 @@ def make_node(self, *inputs):
             batch_shape = tuple(
                 [
                     broadcast_static_dim_lengths(batch_dims)
-                    for batch_dims in zip(*batch_shapes)
+                    for batch_dims in zip(*batch_shapes, strict=True)
                 ]
             )
         except ValueError:
@@ -168,10 +171,10 @@ def infer_shape(
         batch_ndims = self.batch_ndim(node)
         core_dims: dict[str, Any] = {}
         batch_shapes = [input_shape[:batch_ndims] for input_shape in input_shapes]
-        for input_shape, sig in zip(input_shapes, self.inputs_sig):
+        for input_shape, sig in zip(input_shapes, self.inputs_sig, strict=True):
             core_shape = input_shape[batch_ndims:]
 
-            for core_dim, dim_name in zip(core_shape, sig):
+            for core_dim, dim_name in zip(core_shape, sig, strict=True):
                 prev_core_dim = core_dims.get(core_dim)
                 if prev_core_dim is None:
                     core_dims[dim_name] = core_dim
@@ -182,7 +185,7 @@ def infer_shape(
         batch_shape = broadcast_shape(*batch_shapes, arrays_are_shapes=True)
 
         out_shapes = []
-        for output, sig in zip(node.outputs, self.outputs_sig):
+        for output, sig in zip(node.outputs, self.outputs_sig, strict=True):
             core_out_shape = []
             for i, dim_name in enumerate(sig):
                 # The output dim is the same as another input dim
@@ -213,17 +216,17 @@ def as_core(t, core_t):
         with config.change_flags(compute_test_value="off"):
             safe_inputs = [
                 tensor(dtype=inp.type.dtype, shape=(None,) * len(sig))
-                for inp, sig in zip(inputs, self.inputs_sig)
+                for inp, sig in zip(inputs, self.inputs_sig, strict=True)
             ]
             core_node = self._create_dummy_core_node(safe_inputs)
 
             core_inputs = [
                 as_core(inp, core_inp)
-                for inp, core_inp in zip(inputs, core_node.inputs)
+                for inp, core_inp in zip(inputs, core_node.inputs, strict=True)
             ]
             core_ograds = [
                 as_core(ograd, core_ograd)
-                for ograd, core_ograd in zip(ograds, core_node.outputs)
+                for ograd, core_ograd in zip(ograds, core_node.outputs, strict=True)
             ]
             core_outputs = core_node.outputs
 
@@ -232,7 +235,11 @@ def as_core(t, core_t):
         igrads = vectorize_graph(
             [core_igrad for core_igrad in core_igrads if core_igrad is not None],
             replace=dict(
-                zip(core_inputs + core_outputs + core_ograds, inputs + outputs + ograds)
+                zip(
+                    core_inputs + core_outputs + core_ograds,
+                    inputs + outputs + ograds,
+                    strict=True,
+                )
             ),
         )
 
@@ -258,7 +265,7 @@ def L_op(self, inputs, outs, ograds):
             # the return value obviously zero so that gradient.grad can tell
             # this op did the right thing.
             new_rval = []
-            for elem, inp in zip(rval, inputs):
+            for elem, inp in zip(rval, inputs, strict=True):
                 if isinstance(elem.type, NullType | DisconnectedType):
                     new_rval.append(elem)
                 else:
@@ -272,15 +279,17 @@ def L_op(self, inputs, outs, ograds):
         # Sum out the broadcasted dimensions
         batch_ndims = self.batch_ndim(outs[0].owner)
         batch_shape = outs[0].type.shape[:batch_ndims]
-        for i, (inp, sig) in enumerate(zip(inputs, self.inputs_sig)):
+        for i, (inp, sig) in enumerate(zip(inputs, self.inputs_sig, strict=True)):
             if isinstance(rval[i].type, NullType | DisconnectedType):
                 continue
 
             assert inp.type.ndim == batch_ndims + len(sig)
 
             to_sum = [
                 j
-                for j, (inp_s, out_s) in enumerate(zip(inp.type.shape, batch_shape))
+                for j, (inp_s, out_s) in enumerate(
+                    zip(inp.type.shape, batch_shape, strict=False)
+                )
                 if inp_s == 1 and out_s != 1
             ]
             if to_sum:
@@ -320,9 +329,14 @@ def _check_runtime_broadcast(self, node, inputs):
 
         for dims_and_bcast in zip(
             *[
-                zip(input.shape[:batch_ndim], sinput.type.broadcastable[:batch_ndim])
-                for input, sinput in zip(inputs, node.inputs)
-            ]
+                zip(
+                    input.shape[:batch_ndim],
+                    sinput.type.broadcastable[:batch_ndim],
+                    strict=True,
+                )
+                for input, sinput in zip(inputs, node.inputs, strict=True)
+            ],
+            strict=True,
         ):
             if any(d != 1 for d, _ in dims_and_bcast) and (1, False) in dims_and_bcast:
                 raise ValueError(
@@ -343,7 +357,9 @@ def perform(self, node, inputs, output_storage):
         if not isinstance(res, tuple):
             res = (res,)
 
-        for node_out, out_storage, r in zip(node.outputs, output_storage, res):
+        for node_out, out_storage, r in zip(
+            node.outputs, output_storage, res, strict=True
+        ):
             out_dtype = getattr(node_out, "dtype", None)
             if out_dtype and out_dtype != r.dtype:
                 r = np.asarray(r, dtype=out_dtype)

pytensor/tensor/conv/abstract_conv.py

@@ -506,7 +506,7 @@ def check_dim(given, computed):
 
     return all(
         check_dim(given, computed)
-        for (given, computed) in zip(output_shape, computed_output_shape)
+        for (given, computed) in zip(output_shape, computed_output_shape, strict=True)
     )