Make gradient of Fishman sqrt happy

Author: JanLuca

varipeps/ctmrg/projectors.py

@@ -293,7 +293,11 @@ def _split_transfer_fishman(first_tensor, second_tensor, truncation_eps):
         (first_ketbra_S / first_ketbra_S[0]) >= truncation_eps, first_ketbra_S, 0
     )
     first_ketbra_S /= jnp.sum(first_ketbra_S)
-    first_ketbra_S = jnp.sqrt(first_ketbra_S)
+    first_ketbra_S = jnp.where(
+        first_ketbra_S == 0,
+        0,
+        jnp.sqrt(jnp.where(first_ketbra_S == 0, 1, first_ketbra_S)),
+    )
     if first_tensor.ndim == 5:
         first_ketbra_U = first_ketbra_U.reshape(
             first_tensor.shape[0], first_tensor.shape[1], first_ketbra_U.shape[-1]
@@ -332,7 +336,11 @@ def _split_transfer_fishman(first_tensor, second_tensor, truncation_eps):
         (second_ketbra_S / second_ketbra_S[0]) >= truncation_eps, second_ketbra_S, 0
     )
     second_ketbra_S /= jnp.sum(second_ketbra_S)
-    second_ketbra_S = jnp.sqrt(second_ketbra_S)
+    second_ketbra_S = jnp.where(
+        second_ketbra_S == 0,
+        0,
+        jnp.sqrt(jnp.where(second_ketbra_S == 0, 1, second_ketbra_S)),
+    )
     if second_tensor.ndim == 5:
         second_ketbra_U = second_ketbra_U.reshape(
             second_tensor.shape[0],

varipeps/peps/unitcell.py

@@ -495,9 +495,7 @@ def convert_to_split_transfer(
             sanity_checks=False,
         )
 
-    def convert_to_full_transfer(
-        self: T_PEPS_Unit_Cell
-    ) -> T_PEPS_Unit_Cell:
+    def convert_to_full_transfer(self: T_PEPS_Unit_Cell) -> T_PEPS_Unit_Cell:
         """
         Convert the list of unique tensors to the full transfer ansatz.
 

varipeps/utils/svd.py

@@ -55,6 +55,7 @@ def _svd_jvp_rule(primals, tangents):
     ds = jnp.real(jnp.diagonal(dS, 0, -2, -1))
 
     s_diffs = (s_dim + _T(s_dim)) * (s_dim - _T(s_dim))
+    # s_diffs = jnp.where(s_diffs / (s[0] ** 2) >= 1e-12, s_diffs, 0)
     s_diffs_zeros = jnp.ones((), dtype=A.dtype) * (
         s_diffs == 0.0
     )  # is 1. where s_diffs is 0. and is 0. everywhere else