Replace safe_zip by built-in zip function with strict=True

JanLuca · JanLuca · commit 737c22d865ff · 2025-07-23T11:39:14.000+02:00
diff --git a/varipeps/ctmrg/routine.py b/varipeps/ctmrg/routine.py
@@ -771,8 +771,10 @@ def _ctmrg_rev_while_body(carry):
     bar_fixed_point = bar_fixed_point_last_step.replace_unique_tensors(
         [
             t_old.__add__(t_new, checks=False)
-            for t_old, t_new in jax.util.safe_zip(
-                initial_bar.get_unique_tensors(), new_env_bar.get_unique_tensors()
+            for t_old, t_new in zip(
+                initial_bar.get_unique_tensors(),
+                new_env_bar.get_unique_tensors(),
+                strict=True,
             )
         ]
     )
diff --git a/varipeps/mapping/florett_pentagon.py b/varipeps/mapping/florett_pentagon.py
@@ -579,7 +579,7 @@ def _calc_onsite_gate(
     Id_other_sites = jnp.eye(d**7)
 
     for i, (b_e, g_e, blue_e) in enumerate(
-        jax.util.safe_zip(black_gates, green_gates, blue_gates)
+        zip(black_gates, green_gates, blue_gates, strict=True)
     ):
         black_12 = jnp.kron(b_e, Id_other_sites)
 
@@ -673,7 +673,7 @@ def _calc_right_gate(
 
     Id_other_site = jnp.eye(d)
 
-    for i, (b_e, blue_e) in enumerate(jax.util.safe_zip(black_gates, blue_gates)):
+    for i, (b_e, blue_e) in enumerate(zip(black_gates, blue_gates, strict=True)):
         black_51 = jnp.kron(b_e, Id_other_site)
 
         blue_59 = jnp.kron(blue_e, Id_other_site)
@@ -701,7 +701,7 @@ def _calc_down_gate(
 
     Id_other_site = jnp.eye(d**2)
 
-    for i, (b_e, blue_e) in enumerate(jax.util.safe_zip(black_gates, blue_gates)):
+    for i, (b_e, blue_e) in enumerate(zip(black_gates, blue_gates, strict=True)):
         black_91 = jnp.kron(b_e, Id_other_site)
         black_91 = black_91.reshape(d, d, d, d, d, d, d, d)
         black_91 = black_91.transpose(2, 0, 1, 3, 6, 4, 5, 7)
@@ -1157,11 +1157,12 @@ def __call__(
                         )
 
                     for sr_i, (sr_o, sr_h, sr_v, sr_d) in enumerate(
-                        jax.util.safe_zip(
+                        zip(
                             step_result_onsite[: len(self.black_gates)],
                             step_result_horizontal[: len(self.black_gates)],
                             step_result_vertical[: len(self.black_gates)],
                             step_result_diagonal[: len(self.black_gates)],
+                            strict=True,
                         )
                     ):
                         result[sr_i] += sr_o + sr_h + sr_v + sr_d
diff --git a/varipeps/mapping/maple_leaf.py b/varipeps/mapping/maple_leaf.py
@@ -262,7 +262,7 @@ def _calc_onsite_gate(
     single_gates = [None] * result_length
 
     for i, (g_e, b_e, r_e) in enumerate(
-        jax.util.safe_zip(green_gates, blue_gates, red_gates)
+        zip(green_gates, blue_gates, red_gates, strict=True)
     ):
         (
             green_12,
@@ -327,7 +327,7 @@ def _calc_right_gate(
 
     single_gates = [None] * result_length
 
-    for i, (b_e, r_e) in enumerate(jax.util.safe_zip(blue_gates, red_gates)):
+    for i, (b_e, r_e) in enumerate(zip(blue_gates, red_gates, strict=True)):
         red_61, blue_62 = get_right_gates(b_e, r_e, d)
 
         result[i] = red_61 + blue_62
@@ -361,7 +361,7 @@ def _calc_down_gate(
 
     single_gates = [None] * result_length
 
-    for i, (b_e, r_e) in enumerate(jax.util.safe_zip(blue_gates, red_gates)):
+    for i, (b_e, r_e) in enumerate(zip(blue_gates, red_gates, strict=True)):
         blue_35, red_36 = get_down_gates(b_e, r_e, d)
 
         result[i] = blue_35 + red_36
@@ -395,7 +395,7 @@ def _calc_diagonal_gate(
 
     single_gates = [None] * result_length
 
-    for i, (b_e, r_e) in enumerate(jax.util.safe_zip(blue_gates, red_gates)):
+    for i, (b_e, r_e) in enumerate(zip(blue_gates, red_gates, strict=True)):
         blue_41, red_31 = get_diagonal_gates(b_e, r_e, d)
 
         result[i] = blue_41 + red_31
@@ -802,11 +802,12 @@ def __call__(
                         )
 
                     for sr_i, (sr_o, sr_h, sr_v, sr_d) in enumerate(
-                        jax.util.safe_zip(
+                        zip(
                             step_result_onsite[: len(self.green_gates)],
                             step_result_horizontal[: len(self.green_gates)],
                             step_result_vertical[: len(self.green_gates)],
                             step_result_diagonal[: len(self.green_gates)],
+                            strict=True,
                         )
                     ):
                         result[sr_i] += sr_o + sr_h + sr_v + sr_d
diff --git a/varipeps/mapping/square_kagome.py b/varipeps/mapping/square_kagome.py
@@ -769,7 +769,7 @@ def __call__(
                     )
 
                     for sr_i, (sr_h, sr_v) in enumerate(
-                        jax.util.safe_zip(step_result_horizontal, step_result_vertical)
+                        zip(step_result_horizontal, step_result_vertical, strict=True)
                     ):
                         result[sr_i] += sr_h + sr_v
                 elif len(self.cross_gates) > 0:
diff --git a/varipeps/mapping/triangular.py b/varipeps/mapping/triangular.py
@@ -232,9 +232,7 @@ def __call__(
                 )
 
                 for sr_i, (sr_x, sr_y, sr_diagonal) in enumerate(
-                    jax.util.safe_zip(
-                        step_result_x, step_result_y, step_result_diagonal
-                    )
+                    zip(step_result_x, step_result_y, step_result_diagonal, strict=True)
                 ):
                     result[sr_i] += sr_x + sr_y + sr_diagonal
 
@@ -633,7 +631,7 @@ def _calc_quadrat_gate_next_nearest(
     Id_other_sites = jnp.eye(d**2)
 
     for i, (n_e, n_n_e) in enumerate(
-        jax.util.safe_zip(nearest_gates, next_nearest_gates)
+        zip(nearest_gates, next_nearest_gates, strict=True)
     ):
         nearest_34 = jnp.kron(Id_other_sites, n_e)
 
@@ -913,10 +911,11 @@ def __call__(
                 )
 
                 for sr_i, (sr_q, sr_h, sr_v) in enumerate(
-                    jax.util.safe_zip(
+                    zip(
                         step_result_quadrat[: len(self.nearest_neighbor_gates)],
                         step_result_horizontal_rect[: len(self.nearest_neighbor_gates)],
                         step_result_vertical_rect[: len(self.nearest_neighbor_gates)],
+                        strict=True,
                     )
                 ):
                     result[sr_i] += sr_q + sr_h + sr_v
diff --git a/varipeps/optimization/inner_function.py b/varipeps/optimization/inner_function.py
@@ -1,6 +1,5 @@
 import jax.numpy as jnp
 from jax import value_and_grad
-from jax.util import safe_zip
 
 from varipeps import varipeps_config
 from varipeps.peps import PEPS_Unit_Cell
@@ -39,7 +38,7 @@ def _map_tensors(
             old_tensors = unitcell.get_unique_tensors()
             if not all(
                 jnp.allclose(ti, tj_obj.tensor)
-                for ti, tj_obj in safe_zip(peps_tensors, old_tensors)
+                for ti, tj_obj in zip(peps_tensors, old_tensors, strict=True)
             ):
                 raise ValueError(
                     "Input tensors and provided unitcell are not the same state."
@@ -110,14 +109,14 @@ def calc_ctmrg_expectation(
                     spiral_vectors_x = (spiral_vectors_x,)
                 spiral_vectors = tuple(
                     jnp.array((sx, sy))
-                    for sx, sy in safe_zip(spiral_vectors_x, spiral_vectors)
+                    for sx, sy in zip(spiral_vectors_x, spiral_vectors, strict=True)
                 )
             elif spiral_vectors_y is not None:
                 if isinstance(spiral_vectors_y, jnp.ndarray):
                     spiral_vectors_y = (spiral_vectors_y,)
                 spiral_vectors = tuple(
                     jnp.array((sx, sy))
-                    for sx, sy in safe_zip(spiral_vectors, spiral_vectors_y)
+                    for sx, sy in zip(spiral_vectors, spiral_vectors_y, strict=True)
                 )
     else:
         peps_tensors, unitcell = _map_tensors(
@@ -211,14 +210,14 @@ def calc_preconverged_ctmrg_value_and_grad(
                     spiral_vectors_x = (spiral_vectors_x,)
                 spiral_vectors = tuple(
                     jnp.array((sx, sy))
-                    for sx, sy in safe_zip(spiral_vectors_x, spiral_vectors)
+                    for sx, sy in zip(spiral_vectors_x, spiral_vectors, strict=True)
                 )
             elif spiral_vectors_y is not None:
                 if isinstance(spiral_vectors_y, jnp.ndarray):
                     spiral_vectors_y = (spiral_vectors_y,)
                 spiral_vectors = tuple(
                     jnp.array((sx, sy))
-                    for sx, sy in safe_zip(spiral_vectors, spiral_vectors_y)
+                    for sx, sy in zip(spiral_vectors, spiral_vectors_y, strict=True)
                 )
     else:
         peps_tensors, unitcell = _map_tensors(
@@ -293,14 +292,14 @@ def calc_ctmrg_expectation_custom(
                     spiral_vectors_x = (spiral_vectors_x,)
                 spiral_vectors = tuple(
                     jnp.array((sx, sy))
-                    for sx, sy in safe_zip(spiral_vectors_x, spiral_vectors)
+                    for sx, sy in zip(spiral_vectors_x, spiral_vectors, strict=True)
                 )
             elif spiral_vectors_y is not None:
                 if isinstance(spiral_vectors_y, jnp.ndarray):
                     spiral_vectors_y = (spiral_vectors_y,)
                 spiral_vectors = tuple(
                     jnp.array((sx, sy))
-                    for sx, sy in safe_zip(spiral_vectors, spiral_vectors_y)
+                    for sx, sy in zip(spiral_vectors, spiral_vectors_y, strict=True)
                 )
     else:
         peps_tensors, unitcell = _map_tensors(
diff --git a/varipeps/optimization/optimizer.py b/varipeps/optimization/optimizer.py
@@ -13,7 +13,6 @@
 import jax.numpy as jnp
 from jax.lax import scan
 from jax.flatten_util import ravel_pytree
-from jax.util import safe_zip
 
 from varipeps import varipeps_config, varipeps_global_state
 from varipeps.config import Optimizing_Methods
@@ -240,14 +239,14 @@ def _autosave_wrapper(
                     spiral_vectors_x = (spiral_vectors_x,)
                 spiral_vectors = tuple(
                     jnp.array((sx, sy))
-                    for sx, sy in safe_zip(spiral_vectors_x, spiral_vectors)
+                    for sx, sy in zip(spiral_vectors_x, spiral_vectors, strict=True)
                 )
             elif spiral_vectors_y is not None:
                 if isinstance(spiral_vectors_y, jnp.ndarray):
                     spiral_vectors_y = (spiral_vectors_y,)
                 spiral_vectors = tuple(
                     jnp.array((sx, sy))
-                    for sx, sy in safe_zip(spiral_vectors, spiral_vectors_y)
+                    for sx, sy in zip(spiral_vectors, spiral_vectors_y, strict=True)
                 )
     elif additional_input.get("spiral_vectors") is not None:
         spiral_vectors = additional_input.get("spiral_vectors")

Original file line number	Diff line number	Diff line change
`@@ -771,8 +771,10 @@ def _ctmrg_rev_while_body(carry):`
`771`	`771`	`bar_fixed_point = bar_fixed_point_last_step.replace_unique_tensors(`
`772`	`772`	`[`
`773`	`773`	`t_old.__add__(t_new, checks=False)`
`774`		`- for t_old, t_new in jax.util.safe_zip(`
`775`		`- initial_bar.get_unique_tensors(), new_env_bar.get_unique_tensors()`
	`774`	`+ for t_old, t_new in zip(`
	`775`	`+ initial_bar.get_unique_tensors(),`
	`776`	`+ new_env_bar.get_unique_tensors(),`
	`777`	`+ strict=True,`
`776`	`778`	`)`
`777`	`779`	`]`
`778`	`780`	`)`
Original file line number	Diff line number	Diff line change
`@@ -769,7 +769,7 @@ def __call__(`
`769`	`769`	`)`
`770`	`770`
`771`	`771`	`for sr_i, (sr_h, sr_v) in enumerate(`
`772`		`- jax.util.safe_zip(step_result_horizontal, step_result_vertical)`
	`772`	`+ zip(step_result_horizontal, step_result_vertical, strict=True)`
`773`	`773`	`):`
`774`	`774`	`result[sr_i] += sr_h + sr_v`
`775`	`775`	`elif len(self.cross_gates) > 0:`