Try to avoid memory leaks by clearing jax caches if CTMRG heuristics forces a lot of JITs

Author: JanLuca

varipeps/optimization/line_search.py

@@ -2,6 +2,7 @@
 
 from tqdm_loggable.auto import tqdm
 
+import jax
 import jax.numpy as jnp
 from jax import jit
 from jax.flatten_util import ravel_pytree
@@ -1119,6 +1120,9 @@ def line_search(
 
         count += 1
 
+    if new_unitcell[0, 0][0][0].chi != unitcell[0, 0][0][0].chi:
+        jax.clear_caches()
+
     if count == varipeps_config.line_search_max_steps:
         raise NoSuitableStepSizeError(f"Count {count}, Last alpha {alpha}")
 

varipeps/optimization/optimizer.py

@@ -8,6 +8,7 @@
 
 from tqdm_loggable.auto import tqdm
 
+import jax
 from jax import jit
 import jax.numpy as jnp
 from jax.lax import scan
@@ -405,6 +406,7 @@ def random_noise(a):
         while count < varipeps_config.optimizer_max_steps:
             runtime_start = time.perf_counter()
 
+            chi_before_ctmrg = working_unitcell[0, 0][0][0].chi
             try:
                 if varipeps_config.ad_use_custom_vjp:
                     (
@@ -498,6 +500,9 @@ def random_noise(a):
 
                     continue
 
+            if working_unitcell[0, 0][0][0].chi != chi_before_ctmrg:
+                jax.clear_caches()
+
             working_gradient = [elem.conj() for elem in working_gradient_seq]
 
             if signal_reset_descent_dir: