Make zips strict in pytensor/link/numba

Author: Armavica

pytensor/link/numba/dispatch/basic.py

@@ -405,7 +405,7 @@ def py_perform_return(inputs):
         def py_perform_return(inputs):
             return tuple(
                 out_type.filter(out[0])
-                for out_type, out in zip(output_types, py_perform(inputs))
+                for out_type, out in zip(output_types, py_perform(inputs), strict=True)
             )
 
     @numba_njit
@@ -568,7 +568,7 @@ def numba_funcify_SpecifyShape(op, node, **kwargs):
     func_conditions = [
         f"assert x.shape[{i}] == {shape_input_names}"
         for i, (shape_input, shape_input_names) in enumerate(
-            zip(shape_inputs, shape_input_names)
+            zip(shape_inputs, shape_input_names, strict=True)
         )
         if shape_input is not NoneConst
     ]

pytensor/link/numba/dispatch/cython_support.py

@@ -45,7 +45,7 @@ def arg_numba_types(self) -> list[DTypeLike]:
     def can_cast_args(self, args: list[DTypeLike]) -> bool:
         ok = True
         count = 0
-        for name, dtype in zip(self.arg_names, self.arg_dtypes):
+        for name, dtype in zip(self.arg_names, self.arg_dtypes, strict=True):
             if name == "__pyx_skip_dispatch":
                 continue
             if len(args) <= count:
@@ -164,7 +164,10 @@ def __wrapper_address__(self):
         return self._func_ptr
 
     def __call__(self, *args, **kwargs):
-        args = [dtype(arg) for arg, dtype in zip(args, self._signature.arg_dtypes)]
+        args = [
+            dtype(arg)
+            for arg, dtype in zip(args, self._signature.arg_dtypes, strict=True)
+        ]
         if self.has_pyx_skip_dispatch():
             output = self._pyfunc(*args[:-1], **kwargs)
         else:

pytensor/link/numba/dispatch/elemwise.py

@@ -515,8 +515,10 @@ def elemwise(*inputs):
         inputs = [np.asarray(input) for input in inputs]
         inputs_bc = np.broadcast_arrays(*inputs)
         shape = inputs[0].shape
-        for input, bc in zip(inputs, input_bc_patterns):
-            for length, allow_bc, iter_length in zip(input.shape, bc, shape):
+        for input, bc in zip(inputs, input_bc_patterns, strict=True):
+            for length, allow_bc, iter_length in zip(
+                input.shape, bc, shape, strict=True
+            ):
                 if length == 1 and shape and iter_length != 1 and not allow_bc:
                     raise ValueError("Broadcast not allowed.")
 
@@ -529,11 +531,11 @@ def elemwise(*inputs):
             outs = scalar_op_fn(*vals)
             if not isinstance(outs, tuple):
                 outs = (outs,)
-            for out, out_val in zip(outputs, outs):
+            for out, out_val in zip(outputs, outs, strict=True):
                 out[idx] = out_val
 
         outputs_summed = []
-        for output, bc in zip(outputs, output_bc_patterns):
+        for output, bc in zip(outputs, output_bc_patterns, strict=True):
             axes = tuple(np.nonzero(bc)[0])
             outputs_summed.append(output.sum(axes, keepdims=True))
         if len(outputs_summed) != 1:

pytensor/link/numba/dispatch/extra_ops.py

@@ -186,7 +186,7 @@ def ravelmultiindex(*inp):
 
             new_arr = arr.T.astype(np.float64).copy()
             for i, b in enumerate(new_arr):
-                for j, (d, v) in enumerate(zip(shape, b)):
+                for j, (d, v) in enumerate(zip(shape, b)):  # noqa: B905
                     if v < 0 or v >= d:
                         mode_fn(new_arr, i, j, v, d)
 

pytensor/link/numba/dispatch/scalar.py

@@ -118,7 +118,7 @@ def {scalar_op_fn_name}({input_names}):
             [
                 f"direct_cast({i_name}, {i_tmp_dtype_name})"
                 for i_name, i_tmp_dtype_name in zip(
-                    input_names, input_tmp_dtype_names.keys()
+                    input_names, input_tmp_dtype_names.keys(), strict=False
                 )
             ]
         )

pytensor/link/numba/dispatch/scan.py

@@ -163,10 +163,11 @@ def add_inner_in_expr(
             op.info.mit_mot_in_slices
             + op.info.mit_sot_in_slices
             + op.info.sit_sot_in_slices,
+            strict=True,
         )
     )
     inner_in_names_to_output_taps: dict[str, tuple[int, ...] | None] = dict(
-        zip(outer_in_mit_mot_names, op.info.mit_mot_out_slices)
+        zip(outer_in_mit_mot_names, op.info.mit_mot_out_slices, strict=True)
     )
 
     # Inner-outputs consist of:
@@ -373,7 +374,12 @@ def add_output_storage_post_proc_stmt(
     inner_out_post_processing_block = "\n".join(inner_out_post_processing_stmts)
 
     inner_out_to_outer_out_stmts = "\n".join(
-        [f"{s} = {d}" for s, d in zip(inner_out_to_outer_in_stmts, inner_output_names)]
+        [
+            f"{s} = {d}"
+            for s, d in zip(
+                inner_out_to_outer_in_stmts, inner_output_names, strict=True
+            )
+        ]
     )
 
     scan_op_src = f"""

pytensor/link/numba/dispatch/slinalg.py

@@ -421,7 +421,7 @@ def block_diag(*arrs):
         out = np.zeros((out_shape[0], out_shape[1]), dtype=dtype)
 
         r, c = 0, 0
-        for arr, shape in zip(arrs, shapes):
+        for arr, shape in zip(arrs, shapes):  # noqa: B905
             rr, cc = shape
             out[r : r + rr, c : c + cc] = arr
             r += rr

pytensor/link/numba/dispatch/subtensor.py

@@ -158,7 +158,7 @@ def advancedincsubtensor1_inplace(x, val, idxs):
             def advancedincsubtensor1_inplace(x, vals, idxs):
                 if not len(idxs) == len(vals):
                     raise ValueError("The number of indices and values must match.")
-                for idx, val in zip(idxs, vals):
+                for idx, val in zip(idxs, vals):  # noqa: B905
                     x[idx] = val
                 return x
     else:
@@ -184,7 +184,7 @@ def advancedincsubtensor1_inplace(x, val, idxs):
             def advancedincsubtensor1_inplace(x, vals, idxs):
                 if not len(idxs) == len(vals):
                     raise ValueError("The number of indices and values must match.")
-                for idx, val in zip(idxs, vals):
+                for idx, val in zip(idxs, vals):  # noqa: B905
                     x[idx] += val
                 return x
 

pytensor/link/numba/dispatch/tensor_basic.py

@@ -37,7 +37,9 @@ def numba_funcify_AllocEmpty(op, node, **kwargs):
         "\n".join(
             [
                 f"{item_name} = to_scalar({shape_name})"
-                for item_name, shape_name in zip(shape_var_item_names, shape_var_names)
+                for item_name, shape_name in zip(
+                    shape_var_item_names, shape_var_names, strict=True
+                )
             ]
         ),
         " " * 4,
@@ -71,7 +73,9 @@ def numba_funcify_Alloc(op, node, **kwargs):
         "\n".join(
             [
                 f"{item_name} = to_scalar({shape_name})"
-                for item_name, shape_name in zip(shape_var_item_names, shape_var_names)
+                for item_name, shape_name in zip(
+                    shape_var_item_names, shape_var_names, strict=True
+                )
             ]
         ),
         " " * 4,

pytensor/link/numba/dispatch/vectorize_codegen.py

@@ -45,7 +45,7 @@ def store_core_outputs(i0, i1, ..., in, o0, o1, ..., on):
     store_outputs = "\n".join(
         [
             f"{output}[...] = {inner_output}"
-            for output, inner_output in zip(outputs, inner_outputs)
+            for output, inner_output in zip(outputs, inner_outputs, strict=True)
         ]
     )
     func_src = f"""
@@ -139,7 +139,7 @@ def _vectorized(
         )
 
     core_input_types = []
-    for input_type, bc_pattern in zip(input_types, input_bc_patterns):
+    for input_type, bc_pattern in zip(input_types, input_bc_patterns, strict=True):
         core_ndim = input_type.ndim - len(bc_pattern)
         # TODO: Reconsider this
         if core_ndim == 0:
@@ -152,14 +152,18 @@ def _vectorized(
 
     core_out_types = [
         types.Array(numba.from_dtype(np.dtype(dtype)), len(output_core_shape), "C")
-        for dtype, output_core_shape in zip(output_dtypes, output_core_shape_types)
+        for dtype, output_core_shape in zip(
+            output_dtypes, output_core_shape_types, strict=True
+        )
     ]
 
     out_types = [
         types.Array(
             numba.from_dtype(np.dtype(dtype)), batch_ndim + len(output_core_shape), "C"
         )
-        for dtype, output_core_shape in zip(output_dtypes, output_core_shape_types)
+        for dtype, output_core_shape in zip(
+            output_dtypes, output_core_shape_types, strict=True
+        )
     ]
 
     for output_idx, input_idx in inplace_pattern:
@@ -213,7 +217,7 @@ def codegen(
 
         inputs = [
             arrayobj.make_array(ty)(ctx, builder, val)
-            for ty, val in zip(input_types, inputs)
+            for ty, val in zip(input_types, inputs, strict=True)
         ]
         in_shapes = [cgutils.unpack_tuple(builder, obj.shape) for obj in inputs]
 
@@ -285,7 +289,9 @@ def compute_itershape(
     if size is not None:
         shape = size
         for i in range(batch_ndim):
-            for j, (bc, in_shape) in enumerate(zip(broadcast_pattern, in_shapes)):
+            for j, (bc, in_shape) in enumerate(
+                zip(broadcast_pattern, in_shapes, strict=True)
+            ):
                 length = in_shape[i]
                 if bc[i]:
                     with builder.if_then(
@@ -320,7 +326,9 @@ def compute_itershape(
     else:
         # Size is implied by the broadcast pattern
         for i in range(batch_ndim):
-            for j, (bc, in_shape) in enumerate(zip(broadcast_pattern, in_shapes)):
+            for j, (bc, in_shape) in enumerate(
+                zip(broadcast_pattern, in_shapes, strict=True)
+            ):
                 length = in_shape[i]
                 if bc[i]:
                     with builder.if_then(
@@ -376,7 +384,7 @@ def make_outputs(
     one = ir.IntType(64)(1)
     inplace_dict = dict(inplace)
     for i, (core_shape, bc, dtype) in enumerate(
-        zip(output_core_shapes, out_bc, dtypes)
+        zip(output_core_shapes, out_bc, dtypes, strict=True)
     ):
         if i in inplace_dict:
             output_arrays.append(inputs[inplace_dict[i]])
@@ -390,7 +398,8 @@ def make_outputs(
         # This is actually an internal numba function, I guess we could
         # call `numba.nd.unsafe.ndarray` instead?
         batch_shape = [
-            length if not bc_dim else one for length, bc_dim in zip(iter_shape, bc)
+            length if not bc_dim else one
+            for length, bc_dim in zip(iter_shape, bc, strict=True)
         ]
         shape = batch_shape + core_shape
         array = arrayobj._empty_nd_impl(ctx, builder, arrtype, shape)
@@ -460,10 +469,10 @@ def make_loop_call(
 
     # Load values from input arrays
     input_vals = []
-    for input, input_type, bc in zip(inputs, input_types, input_bc):
+    for input, input_type, bc in zip(inputs, input_types, input_bc, strict=True):
         core_ndim = input_type.ndim - len(bc)
 
-        idxs_bc = [zero if bc else idx for idx, bc in zip(idxs, bc)] + [
+        idxs_bc = [zero if bc else idx for idx, bc in zip(idxs, bc, strict=True)] + [
             zero
         ] * core_ndim
         ptr = cgutils.get_item_pointer2(
@@ -508,13 +517,13 @@ def make_loop_call(
 
     # Create output slices to pass to inner func
     output_slices = []
-    for output, output_type, bc in zip(outputs, output_types, output_bc):
+    for output, output_type, bc in zip(outputs, output_types, output_bc, strict=True):
         core_ndim = output_type.ndim - len(bc)
         size_type = output.shape.type.element  # type: ignore
         output_shape = cgutils.unpack_tuple(builder, output.shape)  # type: ignore
         output_strides = cgutils.unpack_tuple(builder, output.strides)  # type: ignore
 
-        idxs_bc = [zero if bc else idx for idx, bc in zip(idxs, bc)] + [
+        idxs_bc = [zero if bc else idx for idx, bc in zip(idxs, bc, strict=True)] + [
             zero
         ] * core_ndim
         ptr = cgutils.get_item_pointer2(