Make zips strict in pytensor/tensor/random

Author: Armavica

pytensor/tensor/random/basic.py

@@ -1862,7 +1862,7 @@ def rng_fn(cls, rng, p, size):
             # to `p.shape[:-1]` in the call to `vsearchsorted` below.
             if len(size) < (p.ndim - 1):
                 raise ValueError("`size` is incompatible with the shape of `p`")
-            for s, ps in zip(reversed(size), reversed(p.shape[:-1])):
+            for s, ps in zip(reversed(size), reversed(p.shape[:-1]), strict=True):
                 if s == 1 and ps != 1:
                     raise ValueError("`size` is incompatible with the shape of `p`")
 

pytensor/tensor/random/op.py

@@ -152,11 +152,13 @@ def _supp_shape_from_params(self, dist_params, param_shapes=None):
 
             # Try to infer missing support dims from signature of params
             for param, param_sig, ndim_params in zip(
-                dist_params, self.inputs_sig, self.ndims_params
+                dist_params, self.inputs_sig, self.ndims_params, strict=True
             ):
                 if ndim_params == 0:
                     continue
-                for param_dim, dim in zip(param.shape[-ndim_params:], param_sig):
+                for param_dim, dim in zip(
+                    param.shape[-ndim_params:], param_sig, strict=True
+                ):
                     if dim in core_out_shape and core_out_shape[dim] is None:
                         core_out_shape[dim] = param_dim
 
@@ -231,7 +233,7 @@ def _infer_shape(
 
             # Fail early when size is incompatible with parameters
             for i, (param, param_ndim_supp) in enumerate(
-                zip(dist_params, self.ndims_params)
+                zip(dist_params, self.ndims_params, strict=True)
             ):
                 param_batched_dims = getattr(param, "ndim", 0) - param_ndim_supp
                 if param_batched_dims > size_len:
@@ -255,7 +257,7 @@ def extract_batch_shape(p, ps, n):
 
             batch_shape = tuple(
                 s if not b else constant(1, "int64")
-                for s, b in zip(shape[:-n], p.type.broadcastable[:-n])
+                for s, b in zip(shape[:-n], p.type.broadcastable[:-n], strict=True)
             )
             return batch_shape
 
@@ -264,7 +266,9 @@ def extract_batch_shape(p, ps, n):
         # independent variate dimensions are left.
         params_batch_shape = tuple(
             extract_batch_shape(p, ps, n)
-            for p, ps, n in zip(dist_params, param_shapes, self.ndims_params)
+            for p, ps, n in zip(
+                dist_params, param_shapes, self.ndims_params, strict=False
+            )
         )
 
         if len(params_batch_shape) == 1:

pytensor/tensor/random/rewriting/basic.py

@@ -172,7 +172,7 @@ def local_dimshuffle_rv_lift(fgraph, node):
 
     # Updates the params to reflect the Dimshuffled dimensions
     new_dist_params = []
-    for param, param_ndim_supp in zip(dist_params, rv_op.ndims_params):
+    for param, param_ndim_supp in zip(dist_params, rv_op.ndims_params, strict=True):
         # Add broadcastable dimensions to the parameters that would have been expanded by the size
         padleft = batched_dims - (param.ndim - param_ndim_supp)
         if padleft > 0:
@@ -290,7 +290,7 @@ def is_nd_advanced_idx(idx, dtype):
     # non-broadcastable (non-degenerate) parameter dims. These parameters and the new size
     # should still correctly broadcast any degenerate parameter dims.
     new_dist_params = []
-    for param, param_ndim_supp in zip(dist_params, rv_op.ndims_params):
+    for param, param_ndim_supp in zip(dist_params, rv_op.ndims_params, strict=True):
         # We first expand any missing parameter dims (and later index them away or keep them with none-slicing)
         batch_param_dims_missing = batch_ndims - (param.ndim - param_ndim_supp)
         batch_param = (
@@ -302,7 +302,7 @@ def is_nd_advanced_idx(idx, dtype):
         bcast_batch_param_dims = tuple(
             dim
             for dim, (param_dim, output_dim) in enumerate(
-                zip(batch_param.type.shape, rv.type.shape)
+                zip(batch_param.type.shape, rv.type.shape, strict=False)
             )
             if (param_dim == 1) and (output_dim != 1)
         )

pytensor/tensor/random/utils.py

@@ -42,7 +42,7 @@ def params_broadcast_shapes(param_shapes, ndims_params, use_pytensor=True):
     max_fn = maximum if use_pytensor else max
 
     rev_extra_dims = []
-    for ndim_param, param_shape in zip(ndims_params, param_shapes):
+    for ndim_param, param_shape in zip(ndims_params, param_shapes, strict=True):
         # We need this in order to use `len`
         param_shape = tuple(param_shape)
         extras = tuple(param_shape[: (len(param_shape) - ndim_param)])
@@ -65,7 +65,7 @@ def max_bcast(x, y):
         (extra_dims + tuple(param_shape)[-ndim_param:])
         if ndim_param > 0
         else extra_dims
-        for ndim_param, param_shape in zip(ndims_params, param_shapes)
+        for ndim_param, param_shape in zip(ndims_params, param_shapes, strict=True)
     ]
 
     return bcast_shapes
@@ -104,7 +104,9 @@ def broadcast_params(params, ndims_params):
     for p in params:
         param_shape = tuple(
             1 if bcast else s
-            for s, bcast in zip(p.shape, getattr(p, "broadcastable", (False,) * p.ndim))
+            for s, bcast in zip(
+                p.shape, getattr(p, "broadcastable", (False,) * p.ndim), strict=True
+            )
         )
         use_pytensor |= isinstance(p, Variable)
         param_shapes.append(param_shape)
@@ -115,7 +117,8 @@ def broadcast_params(params, ndims_params):
     broadcast_to_fn = broadcast_to if use_pytensor else np.broadcast_to
 
     bcast_params = [
-        broadcast_to_fn(param, shape) for shape, param in zip(shapes, params)
+        broadcast_to_fn(param, shape)
+        for shape, param in zip(shapes, params, strict=True)
     ]
 
     return bcast_params
@@ -129,7 +132,8 @@ def explicit_expand_dims(
     """Introduce explicit expand_dims in RV parameters that are implicitly broadcasted together and/or by size."""
 
     batch_dims = [
-        param.type.ndim - ndim_param for param, ndim_param in zip(params, ndim_params)
+        param.type.ndim - ndim_param
+        for param, ndim_param in zip(params, ndim_params, strict=False)
     ]
 
     if size_length is not None:
@@ -138,7 +142,7 @@ def explicit_expand_dims(
         max_batch_dims = max(batch_dims, default=0)
 
     new_params = []
-    for new_param, batch_dim in zip(params, batch_dims):
+    for new_param, batch_dim in zip(params, batch_dims, strict=True):
         missing_dims = max_batch_dims - batch_dim
         if missing_dims:
             new_param = shape_padleft(new_param, missing_dims)
@@ -153,7 +157,7 @@ def compute_batch_shape(
     params = explicit_expand_dims(params, ndims_params)
     batch_params = [
         param[(..., *(0,) * core_ndim)]
-        for param, core_ndim in zip(params, ndims_params)
+        for param, core_ndim in zip(params, ndims_params, strict=True)
     ]
     return broadcast_arrays(*batch_params)[0].shape
 
@@ -269,7 +273,9 @@ def seed(self, seed=None):
         self.gen_seedgen = np.random.SeedSequence(seed)
         old_r_seeds = self.gen_seedgen.spawn(len(self.state_updates))
 
-        for (old_r, new_r), old_r_seed in zip(self.state_updates, old_r_seeds):
+        for (old_r, new_r), old_r_seed in zip(
+            self.state_updates, old_r_seeds, strict=True
+        ):
             old_r.set_value(self.rng_ctor(old_r_seed), borrow=True)
 
     def gen(self, op: "RandomVariable", *args, **kwargs) -> TensorVariable: