Make zips strict in pytensor/link/jax

Author: Armavica

pytensor/link/jax/dispatch/scan.py

@@ -30,7 +30,9 @@ def scan(*outer_inputs):
         seqs = op.outer_seqs(outer_inputs)  # JAX `xs`
 
         mit_sot_init = []
-        for tap, seq in zip(op.info.mit_sot_in_slices, op.outer_mitsot(outer_inputs)):
+        for tap, seq in zip(
+            op.info.mit_sot_in_slices, op.outer_mitsot(outer_inputs), strict=True
+        ):
             init_slice = seq[: abs(min(tap))]
             mit_sot_init.append(init_slice)
 
@@ -61,7 +63,9 @@ def jax_args_to_inner_func_args(carry, x):
             inner_seqs = x
 
             mit_sot_flatten = []
-            for array, index in zip(inner_mit_sot, op.info.mit_sot_in_slices):
+            for array, index in zip(
+                inner_mit_sot, op.info.mit_sot_in_slices, strict=True
+            ):
                 mit_sot_flatten.extend(array[jnp.array(index)])
 
             inner_scan_inputs = [
@@ -98,8 +102,7 @@ def inner_func_outs_to_jax_outs(
             inner_mit_sot_new = [
                 jnp.concatenate([old_mit_sot[1:], new_val[None, ...]], axis=0)
                 for old_mit_sot, new_val in zip(
-                    inner_mit_sot,
-                    inner_mit_sot_outs,
+                    inner_mit_sot, inner_mit_sot_outs, strict=True
                 )
             ]
 
@@ -152,7 +155,9 @@ def get_partial_traces(traces):
                 + op.outer_nitsot(outer_inputs)
             )
             partial_traces = []
-            for init_state, trace, buffer in zip(init_states, traces, buffers):
+            for init_state, trace, buffer in zip(
+                init_states, traces, buffers, strict=True
+            ):
                 if init_state is not None:
                     # MIT-SOT and SIT-SOT: The final output should be as long as the input buffer
                     trace = jnp.atleast_1d(trace)

pytensor/link/jax/dispatch/shape.py

@@ -96,7 +96,7 @@ def shape_i(x):
 def jax_funcify_SpecifyShape(op, node, **kwargs):
     def specifyshape(x, *shape):
         assert x.ndim == len(shape)
-        for actual, expected in zip(x.shape, shape):
+        for actual, expected in zip(x.shape, shape, strict=True):
             if expected is None:
                 continue
             if actual != expected:

pytensor/link/jax/dispatch/tensor_basic.py

@@ -200,7 +200,8 @@ def jax_funcify_Tri(op, node, **kwargs):
     def tri(*args):
         # args is N, M, k
         args = [
-            x if const_x is None else const_x for x, const_x in zip(args, const_args)
+            x if const_x is None else const_x
+            for x, const_x in zip(args, const_args, strict=True)
         ]
         return jnp.tri(*args, dtype=op.dtype)
 

pytensor/link/jax/linker.py

@@ -35,12 +35,14 @@ def fgraph_convert(self, fgraph, input_storage, storage_map, **kwargs):
             ]
 
             fgraph.replace_all(
-                zip(shared_rng_inputs, new_shared_rng_inputs),
+                zip(shared_rng_inputs, new_shared_rng_inputs, strict=True),
                 import_missing=True,
                 reason="JAXLinker.fgraph_convert",
             )
 
-            for old_inp, new_inp in zip(shared_rng_inputs, new_shared_rng_inputs):
+            for old_inp, new_inp in zip(
+                shared_rng_inputs, new_shared_rng_inputs, strict=True
+            ):
                 new_inp_storage = [new_inp.get_value(borrow=True)]
                 storage_map[new_inp] = new_inp_storage
                 old_inp_storage = storage_map.pop(old_inp)