Replace safe_zip by built-in zip function for triangular CTMRG

JanLuca · JanLuca · commit 35c2b0cf8ffc · 2025-09-17T15:15:25.000+02:00
diff --git a/varipeps/expectation/triangular_next_nearest.py b/varipeps/expectation/triangular_next_nearest.py
@@ -544,13 +544,14 @@ def __call__(
                 )
 
                 for sr_i, (sr_h, sr_v, sr_d, sr_np, sr_p2p, sr_2pp) in enumerate(
-                    jax.util.safe_zip(
+                    zip(
                         step_result_horizontal,
                         step_result_vertical,
                         step_result_diagonal,
                         step_result_nn_neg_pos,
                         step_result_nn_pos_2pos,
                         step_result_nn_2pos_pos,
+                        strict=True,
                     )
                 ):
                     result[sr_i] += sr_h + sr_v + sr_d + sr_np + sr_p2p + sr_2pp
diff --git a/varipeps/expectation/triangular_two_sites.py b/varipeps/expectation/triangular_two_sites.py
@@ -281,10 +281,11 @@ def __call__(
                 )
 
                 for sr_i, (sr_h, sr_v, sr_d) in enumerate(
-                    jax.util.safe_zip(
+                    zip(
                         step_result_horizontal,
                         step_result_vertical,
                         step_result_diagonal,
+                        strict=True,
                     )
                 ):
                     result[sr_i] += sr_h + sr_v + sr_d
diff --git a/varipeps/mapping/kagome.py b/varipeps/mapping/kagome.py
@@ -2007,11 +2007,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.up_nearest_gates)],
                             step_result_horizontal[: len(self.up_nearest_gates)],
                             step_result_vertical[: len(self.up_nearest_gates)],
                             step_result_diagonal[: len(self.up_nearest_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
@@ -2238,11 +2238,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

Original file line number	Diff line number	Diff line change
`@@ -281,10 +281,11 @@ def __call__(`
`281`	`281`	`)`
`282`	`282`
`283`	`283`	`for sr_i, (sr_h, sr_v, sr_d) in enumerate(`
`284`		`- jax.util.safe_zip(`
	`284`	`+ zip(`
`285`	`285`	`step_result_horizontal,`
`286`	`286`	`step_result_vertical,`
`287`	`287`	`step_result_diagonal,`
	`288`	`+ strict=True,`
`288`	`289`	`)`
`289`	`290`	`):`
`290`	`291`	`result[sr_i] += sr_h + sr_v + sr_d`
Original file line number	Diff line number	Diff line change
`@@ -2007,11 +2007,12 @@ def __call__(`
`2007`	`2007`	`)`
`2008`	`2008`
`2009`	`2009`	`for sr_i, (sr_o, sr_h, sr_v, sr_d) in enumerate(`
`2010`		`- jax.util.safe_zip(`
	`2010`	`+ zip(`
`2011`	`2011`	`step_result_onsite[: len(self.up_nearest_gates)],`
`2012`	`2012`	`step_result_horizontal[: len(self.up_nearest_gates)],`
`2013`	`2013`	`step_result_vertical[: len(self.up_nearest_gates)],`
`2014`	`2014`	`step_result_diagonal[: len(self.up_nearest_gates)],`
	`2015`	`+ strict=True,`
`2015`	`2016`	`)`
`2016`	`2017`	`):`
`2017`	`2018`	`result[sr_i] += sr_o + sr_h + sr_v + sr_d`
Original file line number	Diff line number	Diff line change
`@@ -2238,11 +2238,12 @@ def __call__(`
`2238`	`2238`	`)`
`2239`	`2239`
`2240`	`2240`	`for sr_i, (sr_o, sr_h, sr_v, sr_d) in enumerate(`
`2241`		`- jax.util.safe_zip(`
	`2241`	`+ zip(`
`2242`	`2242`	`step_result_onsite[: len(self.green_gates)],`
`2243`	`2243`	`step_result_horizontal[: len(self.green_gates)],`
`2244`	`2244`	`step_result_vertical[: len(self.green_gates)],`
`2245`	`2245`	`step_result_diagonal[: len(self.green_gates)],`
	`2246`	`+ strict=True,`
`2246`	`2247`	`)`
`2247`	`2248`	`):`
`2248`	`2249`	`result[sr_i] += sr_o + sr_h + sr_v + sr_d`