Make all zips strict in tests

Author: Armavica

tests/compile/function/test_types.py

@@ -371,7 +371,7 @@ def test_copy_share_memory(self):
 
             # Assert storages of SharedVariable without updates are shared
             for (input, _1, _2), here, there in zip(
-                ori.indices, ori.input_storage, cpy.input_storage
+                ori.indices, ori.input_storage, cpy.input_storage, strict=True
             ):
                 assert here.data is there.data
 
@@ -467,7 +467,7 @@ def test_swap_SharedVariable_with_given(self):
             swap={train_x: test_x, train_y: test_y}, delete_updates=True
         )
 
-        for in1, in2 in zip(test_def.maker.inputs, test_cpy.maker.inputs):
+        for in1, in2 in zip(test_def.maker.inputs, test_cpy.maker.inputs, strict=True):
             assert in1.value is in2.value
 
     def test_copy_delete_updates(self):
@@ -899,7 +899,7 @@ def test_deepcopy(self):
         # print 'f.defaults = %s' % (f.defaults, )
         # print 'g.defaults = %s' % (g.defaults, )
         for (f_req, f_feed, f_val), (g_req, g_feed, g_val) in zip(
-            f.defaults, g.defaults
+            f.defaults, g.defaults, strict=True
         ):
             assert f_req == g_req and f_feed == g_feed and f_val == g_val
 
@@ -1070,7 +1070,7 @@ def test_optimizations_preserved(self):
         tf = f.maker.fgraph.toposort()
         tg = f.maker.fgraph.toposort()
         assert len(tf) == len(tg)
-        for nf, ng in zip(tf, tg):
+        for nf, ng in zip(tf, tg, strict=True):
             assert nf.op == ng.op
             assert len(nf.inputs) == len(ng.inputs)
             assert len(nf.outputs) == len(ng.outputs)

tests/compile/test_builders.py

@@ -722,5 +722,5 @@ def test_debugprint():
        └─ *2-<Matrix(float64, shape=(?, ?))> [id I]
 """
 
-    for truth, out in zip(exp_res.split("\n"), lines):
+    for truth, out in zip(exp_res.split("\n"), lines, strict=True):
         assert truth.strip() == out.strip()

tests/d3viz/test_formatting.py

@@ -19,7 +19,7 @@ def setup_method(self):
     def node_counts(self, graph):
         node_types = [node.get_attributes()["node_type"] for node in graph.get_nodes()]
         a, b = np.unique(node_types, return_counts=True)
-        nc = dict(zip(a, b))
+        nc = dict(zip(a, b, strict=True))
         return nc
 
     @pytest.mark.parametrize("mode", ["FAST_RUN", "FAST_COMPILE"])

tests/graph/test_fg.py

@@ -31,13 +31,22 @@ def test_pickle(self):
         s = pickle.dumps(func)
         new_func = pickle.loads(s)
 
-        assert all(type(a) is type(b) for a, b in zip(func.inputs, new_func.inputs))
-        assert all(type(a) is type(b) for a, b in zip(func.outputs, new_func.outputs))
+        assert all(
+            type(a) is type(b)
+            for a, b in zip(func.inputs, new_func.inputs, strict=True)
+        )
+        assert all(
+            type(a) is type(b)
+            for a, b in zip(func.outputs, new_func.outputs, strict=True)
+        )
         assert all(
             type(a.op) is type(b.op)  # noqa: E721
-            for a, b in zip(func.apply_nodes, new_func.apply_nodes)
+            for a, b in zip(func.apply_nodes, new_func.apply_nodes, strict=True)
+        )
+        assert all(
+            a.type == b.type
+            for a, b in zip(func.variables, new_func.variables, strict=True)
         )
-        assert all(a.type == b.type for a, b in zip(func.variables, new_func.variables))
 
     def test_validate_inputs(self):
         var1 = op1()

tests/graph/utils.py

@@ -137,7 +137,9 @@ def __init__(self, inner_inputs, inner_outputs):
             if not isinstance(v, Constant)
         ]
         outputs = clone_replace(inner_outputs, replace=input_replacements)
-        _, inputs = zip(*input_replacements) if input_replacements else (None, [])
+        _, inputs = (
+            zip(*input_replacements, strict=True) if input_replacements else (None, [])
+        )
         self.fgraph = FunctionGraph(inputs, outputs, clone=False)
 
     def make_node(self, *inputs):

tests/link/jax/test_basic.py

@@ -79,7 +79,7 @@ def compare_jax_and_py(
     py_res = pytensor_py_fn(*test_inputs)
 
     if len(fgraph.outputs) > 1:
-        for j, p in zip(jax_res, py_res):
+        for j, p in zip(jax_res, py_res, strict=True):
             assert_fn(j, p)
     else:
         assert_fn(jax_res, py_res)

tests/link/jax/test_random.py

@@ -61,7 +61,9 @@ def test_random_updates(rng_ctor):
     # Check that original rng variable content was not overwritten when calling jax_typify
     assert all(
         a == b if not isinstance(a, np.ndarray) else np.array_equal(a, b)
-        for a, b in zip(rng.get_value().__getstate__(), original_value.__getstate__())
+        for a, b in zip(
+            rng.get_value().__getstate__(), original_value.__getstate__(), strict=True
+        )
     )
 
 
@@ -92,7 +94,9 @@ def test_replaced_shared_rng_storage_order(noise_first):
     ), "Test may need to be tweaked"
 
     # Confirm that input_storage type and fgraph input order are aligned
-    for storage, fgrapn_input in zip(f.input_storage, f.maker.fgraph.inputs):
+    for storage, fgrapn_input in zip(
+        f.input_storage, f.maker.fgraph.inputs, strict=True
+    ):
         assert storage.type == fgrapn_input.type
 
     assert mu.get_value() == 1

tests/link/numba/test_basic.py

@@ -292,7 +292,7 @@ def assert_fn(x, y):
         eval_python_only(fn_inputs, fn_outputs, inputs, mode=numba_mode)
 
     if len(fn_outputs) > 1:
-        for j, p in zip(numba_res, py_res):
+        for j, p in zip(numba_res, py_res, strict=True):
             assert_fn(j, p)
     else:
         assert_fn(numba_res[0], py_res[0])

tests/link/numba/test_scan.py

@@ -488,7 +488,7 @@ def step(seq1, seq2, mitsot1, mitsot2, sitsot1):
 
     ref_fn = pytensor.function(list(test.keys()), outs, mode=get_mode("FAST_COMPILE"))
     ref_res = ref_fn(*test.values())
-    for numba_r, ref_r in zip(numba_res, ref_res):
+    for numba_r, ref_r in zip(numba_res, ref_res, strict=True):
         np.testing.assert_array_almost_equal(numba_r, ref_r)
 
     benchmark(numba_fn, *test.values())

tests/link/pytorch/test_basic.py

@@ -65,7 +65,7 @@ def compare_pytorch_and_py(
     py_res = pytensor_py_fn(*test_inputs)
 
     if len(fgraph.outputs) > 1:
-        for j, p in zip(pytorch_res, py_res):
+        for j, p in zip(pytorch_res, py_res, strict=True):
             assert_fn(j.cpu(), p)
     else:
         assert_fn([pytorch_res[0].cpu()], py_res)