Update dependencies and adapt code to new jax version

Author: JanLuca

peps_ad/__init__.py

@@ -14,7 +14,7 @@
 
 from .__version__ import __version__
 
-from jax.config import config as jax_config
+from jax import config as jax_config
 
 jax_config.update("jax_enable_x64", True)
 

peps_ad/ctmrg/projectors.py

@@ -349,12 +349,16 @@ def calc_left_projectors(
         top_right,
         bottom_left,
         bottom_right,
-        config.ctmrg_truncation_eps
-        if state.ctmrg_effective_truncation_eps is None
-        else state.ctmrg_effective_truncation_eps,
-        config.ctmrg_full_projector_method
-        if state.ctmrg_projector_method is None
-        else state.ctmrg_projector_method,
+        (
+            config.ctmrg_truncation_eps
+            if state.ctmrg_effective_truncation_eps is None
+            else state.ctmrg_effective_truncation_eps
+        ),
+        (
+            config.ctmrg_full_projector_method
+            if state.ctmrg_projector_method is None
+            else state.ctmrg_projector_method
+        ),
         chi,
     )
 
@@ -456,12 +460,16 @@ def calc_right_projectors(
         top_right,
         bottom_left,
         bottom_right,
-        config.ctmrg_truncation_eps
-        if state.ctmrg_effective_truncation_eps is None
-        else state.ctmrg_effective_truncation_eps,
-        config.ctmrg_full_projector_method
-        if state.ctmrg_projector_method is None
-        else state.ctmrg_projector_method,
+        (
+            config.ctmrg_truncation_eps
+            if state.ctmrg_effective_truncation_eps is None
+            else state.ctmrg_effective_truncation_eps
+        ),
+        (
+            config.ctmrg_full_projector_method
+            if state.ctmrg_projector_method is None
+            else state.ctmrg_projector_method
+        ),
         chi,
     )
 
@@ -563,12 +571,16 @@ def calc_top_projectors(
         top_right,
         bottom_left,
         bottom_right,
-        config.ctmrg_truncation_eps
-        if state.ctmrg_effective_truncation_eps is None
-        else state.ctmrg_effective_truncation_eps,
-        config.ctmrg_full_projector_method
-        if state.ctmrg_projector_method is None
-        else state.ctmrg_projector_method,
+        (
+            config.ctmrg_truncation_eps
+            if state.ctmrg_effective_truncation_eps is None
+            else state.ctmrg_effective_truncation_eps
+        ),
+        (
+            config.ctmrg_full_projector_method
+            if state.ctmrg_projector_method is None
+            else state.ctmrg_projector_method
+        ),
         chi,
     )
 
@@ -670,11 +682,15 @@ def calc_bottom_projectors(
         top_right,
         bottom_left,
         bottom_right,
-        config.ctmrg_truncation_eps
-        if state.ctmrg_effective_truncation_eps is None
-        else state.ctmrg_effective_truncation_eps,
-        config.ctmrg_full_projector_method
-        if state.ctmrg_projector_method is None
-        else state.ctmrg_projector_method,
+        (
+            config.ctmrg_truncation_eps
+            if state.ctmrg_effective_truncation_eps is None
+            else state.ctmrg_effective_truncation_eps
+        ),
+        (
+            config.ctmrg_full_projector_method
+            if state.ctmrg_projector_method is None
+            else state.ctmrg_projector_method
+        ),
         chi,
     )

peps_ad/ctmrg/routine.py

@@ -413,9 +413,11 @@ def calc_ctmrg_env(
                 peps_tensors,
                 working_unitcell,
                 False,
-                corner_singular_vals
-                if corner_singular_vals is not None
-                else init_corner_singular_vals,
+                (
+                    corner_singular_vals
+                    if corner_singular_vals is not None
+                    else init_corner_singular_vals
+                ),
                 eps,
                 tmp_count,
                 enforce_elementwise_convergence,

peps_ad/optimization/basinhopping.py

@@ -50,9 +50,9 @@ class PEPS_AD_Basinhopping:
     expectation_func: Expectation_Model
     convert_to_unitcell_func: Optional[Map_To_PEPS_Model] = None
     autosave_filename: PathLike = "data/autosave.hdf5"
-    autosave_func: Callable[
-        [PathLike, Sequence[jnp.ndarray], PEPS_Unit_Cell], None
-    ] = autosave_function
+    autosave_func: Callable[[PathLike, Sequence[jnp.ndarray], PEPS_Unit_Cell], None] = (
+        autosave_function
+    )
 
     def __post_init__(self):
         if isinstance(self.initial_guess, PEPS_Unit_Cell):

peps_ad/optimization/optimizer.py

@@ -300,9 +300,9 @@ def random_noise(a):
         l_bfgs_grad_cache = deque(maxlen=peps_ad_config.optimizer_l_bfgs_maxlen + 1)
 
     count = 0
-    linesearch_step: Union[
-        float, jnp.ndarray
-    ] = peps_ad_config.line_search_initial_step_size
+    linesearch_step: Union[float, jnp.ndarray] = (
+        peps_ad_config.line_search_initial_step_size
+    )
     working_value: Union[float, jnp.ndarray]
     max_trunc_error_list = []
 

peps_ad/peps/unitcell.py

@@ -891,10 +891,10 @@ def load_from_group(
                 if grp["config"].attrs.get(config_attr) is not None
             }
             if config_dict.get("ctmrg_full_projector_method"):
-                config_dict[
-                    "ctmrg_full_projector_method"
-                ] = peps_ad.config.Projector_Method(
-                    config_dict["ctmrg_full_projector_method"]
+                config_dict["ctmrg_full_projector_method"] = (
+                    peps_ad.config.Projector_Method(
+                        config_dict["ctmrg_full_projector_method"]
+                    )
                 )
             if config_dict.get("optimizer_method"):
                 config_dict["optimizer_method"] = peps_ad.config.Optimizing_Methods(

peps_ad/utils/svd.py

@@ -2,8 +2,11 @@
 
 import jax.numpy as jnp
 from jax.lax import scan
+from jax.lax.linalg import svd as lax_svd
 from jax import jit, custom_jvp, lax
 
+from jax._src.numpy.util import promote_dtypes_inexact, check_arraylike
+
 from peps_ad import peps_ad_config
 
 from typing import Tuple
@@ -19,18 +22,20 @@ def _H(x):
 
 @custom_jvp
 def svd_wrapper(a):
-    return jnp.linalg.svd(a, full_matrices=False, compute_uv=True)
+    check_arraylike("jnp.linalg.svd", a)
+    (a,) = promote_dtypes_inexact(jnp.asarray(a))
+    return lax_svd(a, full_matrices=False, compute_uv=True)
 
 
 @svd_wrapper.defjvp
 def _svd_jvp_rule(primals, tangents):
     (A,) = primals
     (dA,) = tangents
-    U, s, Vt = jnp.linalg.svd(A, full_matrices=False, compute_uv=True)
+    U, s, Vt = lax_svd(A, full_matrices=False, compute_uv=True)
 
     Ut, V = _H(U), _H(Vt)
     s_dim = s[..., None, :]
-    dS = jnp.matmul(jnp.matmul(Ut, dA), V)
+    dS = Ut @ dA @ V
     ds = jnp.real(jnp.diagonal(dS, 0, -2, -1))
 
     s_diffs = (s_dim + _T(s_dim)) * (s_dim - _T(s_dim))
@@ -39,23 +44,23 @@ def _svd_jvp_rule(primals, tangents):
     )  # is 1. where s_diffs is 0. and is 0. everywhere else
     s_diffs_zeros = lax.expand_dims(s_diffs_zeros, range(s_diffs.ndim - 2))
     F = 1 / (s_diffs + s_diffs_zeros) - s_diffs_zeros
-    dSS = s_dim * dS  # dS.dot(jnp.diag(s))
-    SdS = _T(s_dim) * dS  # jnp.diag(s).dot(dS)
+    dSS = s_dim.astype(A.dtype) * dS  # dS.dot(jnp.diag(s))
+    SdS = _T(s_dim.astype(A.dtype)) * dS  # jnp.diag(s).dot(dS)
 
-    s_zeros = jnp.ones((), dtype=A.dtype) * (s == 0.0)
+    s_zeros = (s == 0).astype(s.dtype)
     s_inv = 1 / (s + s_zeros) - s_zeros
     s_inv_mat = jnp.vectorize(jnp.diag, signature="(k)->(k,k)")(s_inv)
-    dUdV_diag = 0.5 * (dS - _H(dS)) * s_inv_mat
-    dU = jnp.matmul(U, F * (dSS + _H(dSS)) + dUdV_diag)
-    dV = jnp.matmul(V, F * (SdS + _H(SdS)))
+    dUdV_diag = 0.5 * (dS - _H(dS)) * s_inv_mat.astype(A.dtype)
+    dU = U @ (F.astype(A.dtype) * (dSS + _H(dSS)) + dUdV_diag)
+    dV = V @ (F.astype(A.dtype) * (SdS + _H(SdS)))
 
     m, n = A.shape[-2:]
     if m > n:
-        I = lax.expand_dims(jnp.eye(m, dtype=A.dtype), range(U.ndim - 2))
-        dU = dU + jnp.matmul(I - jnp.matmul(U, Ut), jnp.matmul(dA, V)) / s_dim
+        dAV = dA @ V
+        dU = dU + (dAV - U @ (Ut @ dAV)) / s_dim.astype(A.dtype)
     if n > m:
-        I = lax.expand_dims(jnp.eye(n, dtype=A.dtype), range(V.ndim - 2))
-        dV = dV + jnp.matmul(I - jnp.matmul(V, Vt), jnp.matmul(_H(dA), U)) / s_dim
+        dAHU = _H(dA) @ U
+        dV = dV + (dAHU - V @ (Vt @ dAHU)) / s_dim.astype(A.dtype)
 
     return (U, s, Vt), (dU, ds, _H(dV))