Repository: google/TensorNetwork

Issue:
fully jitted eigs for JAX backend
Current implementation is not fully jittable


Discussion:

Resolving patch:
diff --git a/tensornetwork/backends/jax/jax_backend.py b/tensornetwork/backends/jax/jax_backend.py
--- a/tensornetwork/backends/jax/jax_backend.py
+++ b/tensornetwork/backends/jax/jax_backend.py
@@ -26,14 +26,6 @@
 
 _CACHED_MATVECS = {}
 _CACHED_FUNCTIONS = {}
-_MIN_RES_THRESHS = {
-    np.dtype(np.float16): 1E-3,
-    np.dtype(np.float32): 1E-6,
-    np.dtype(np.float64): 1E-12,
-    np.dtype(np.complex128): 1E-12,
-    np.dtype(np.complex64): 1E-6
-}
-
 
 class JaxBackend(abstract_backend.AbstractBackend):
   """See abstract_backend.AbstractBackend for documentation."""
@@ -247,8 +239,7 @@ def eigs(self, #pylint: disable=arguments-differ
            numeig: int = 6,
            tol: float = 1E-8,
            which: Text = 'LR',
-           maxiter: int = 20,
-           res_thresh: Optional[float] = None) -> Tuple[Tensor, List]:
+           maxiter: int = 20) -> Tuple[Tensor, List]:
     """
     Implicitly restarted Arnoldi method for finding the lowest
     eigenvector-eigenvalue pairs of a linear operator `A`.
@@ -308,10 +299,6 @@ def A(H,x):
         (larges magnitude).
       maxiter: Maximum number of restarts. For `maxiter=0` the routine becomes
         equivalent to a simple Arnoldi method.
-      res_thresh: Threshold parameter. Implicitly restarted arnoldi terminates
-        if the norm of the residual `fk` of the shifted arnoldi factorization
-        falls below `res_thresh`. If `None` a default value depending on the
-        `dtype` of the operator is chosen.
     Returns:
       (eigvals, eigvecs)
        eigvals: A list of `numeig` eigenvalues
@@ -336,22 +323,26 @@ def A(H,x):
       raise TypeError("Expected a `jax.array`. Got {}".format(
           type(initial_state)))
 
-    if res_thresh is None:
-      try:
-        res_thresh = _MIN_RES_THRESHS[initial_state.dtype]
-      except KeyError as err:
-        raise KeyError(f"dtype {initial_state.dtype} not supported") from err
     if A not in _CACHED_MATVECS:
       _CACHED_MATVECS[A] = libjax.tree_util.Partial(libjax.jit(A))
 
     if "imp_arnoldi" not in _CACHED_FUNCTIONS:
       imp_arnoldi = jitted_functions._implicitly_restarted_arnoldi(libjax)
       _CACHED_FUNCTIONS["imp_arnoldi"] = imp_arnoldi
-    return _CACHED_FUNCTIONS["imp_arnoldi"](_CACHED_MATVECS[A], args,
-                                            initial_state, num_krylov_vecs,
-                                            numeig, which, tol, maxiter,
-                                            res_thresh, self.jax_precision)
 
+    eta, U, numits = _CACHED_FUNCTIONS["imp_arnoldi"](_CACHED_MATVECS[A], args,
+                                                      initial_state,
+                                                      num_krylov_vecs, numeig,
+                                                      which, tol, maxiter,
+                                                      self.jax_precision)
+    if numeig > numits:
+      warnings.warn(
+          f"Arnoldi terminated early after numits = {numits}"
+          f" < numeig = {numeig} steps. For this value of `numeig `"
+          f"the routine will return spurious eigenvalues of value 0.0."
+          f"Use a smaller value of numeig, or a smaller value for `tol`")
+    return eta, U
+  
   def eigsh_lanczos(
       self,
       A: Callable,
diff --git a/tensornetwork/backends/jax/jitted_functions.py b/tensornetwork/backends/jax/jitted_functions.py
--- a/tensornetwork/backends/jax/jitted_functions.py
+++ b/tensornetwork/backends/jax/jitted_functions.py
@@ -1,5 +1,5 @@
 import functools
-from typing import List, Any, Tuple, Callable, Sequence
+from typing import List, Any, Tuple, Callable, Sequence, Text
 import collections
 import types
 import numpy as np
@@ -253,7 +253,8 @@ def _arnoldi_fact(
       matvec: Callable, args: List, v0: jax.ShapedArray,
       krylov_vectors: jax.ShapedArray, H: jax.ShapedArray, start: int,
       num_krylov_vecs: int, eps: float, precision: JaxPrecisionType
-  ) -> Tuple[jax.ShapedArray, jax.ShapedArray, int]:
+  ) -> Tuple[jax.ShapedArray, jax.ShapedArray, jax.ShapedArray, float, int,
+             bool]:
     """
     Compute an m-step arnoldi factorization of `matvec`, with
     m = min(`it`,`num_krylov_vecs`). The factorization will
@@ -288,23 +289,31 @@ def _arnoldi_fact(
       eps: Convergence parameter. Iteration is terminated if the norm of a
         krylov-vector falls below `eps`.
     Returns:
-      kv: An array of krylov vectors
-      H: A matrix of overlaps
-      it: The number of performed iterations.
+      jax.ShapedArray: An array of shape 
+        `(num_krylov_vecs, np.prod(initial_state.shape))` of krylov vectors.
+      jax.ShapedArray: Upper Hessenberg matrix of shape 
+        `(num_krylov_vecs, num_krylov_vecs`) of the Arnoldi processs.
+      jax.ShapedArray: The unnormalized residual of the Arnoldi process.
+      int: The norm of the residual.
+      int: The number of performed iterations.
+      bool: if `True`: iteration hit an invariant subspace.
+            if `False`: iteration terminated without encountering
+            an invariant subspace.
     """
-    Z = jax.numpy.linalg.norm(v0)
-    v = v0 / Z
-    krylov_vectors = jax.ops.index_update(krylov_vectors,
-                                          jax.ops.index[start, :],
-                                          jax.numpy.ravel(v))
-    H = jax.lax.cond(
-        start > 0, start,
-        lambda x: jax.ops.index_update(H, jax.ops.index[x, x - 1], Z), None,
-        lambda x: H)
 
-    def modified_gram_schmidt_step_arnoldi(j, vals):
+    # Note (mganahl): currently unused, but is very convenient to have
+    # for further development and tests (it's usually more accurate than
+    # classical gs)
+    # Call signature:
+    #```python
+    # initial_vals = [Av.ravel(), krylov_vectors, i, H]
+    # Av, krylov_vectors, _, H = jax.lax.fori_loop(
+    #     0, i + 1, modified_gram_schmidt_step_arnoldi, initial_vals)
+    #```
+    def modified_gram_schmidt_step_arnoldi(j, vals): #pylint: disable=unused-variable
       """
       Single step of a modified gram-schmidt orthogonalization.
+      Substantially more accurate than classical gram schmidt
       Args:
         j: Integer value denoting the vector to be orthogonalized.
         vals: A list of variables:
@@ -321,44 +330,85 @@ def modified_gram_schmidt_step_arnoldi(j, vals):
       v = krylov_vectors[j, :]
       h = jax.numpy.vdot(v, vector, precision=precision)
       H = H.at[j, n].set(h)
-      vector = vector - h * jax.numpy.reshape(v, vector.shape)
+      vector = vector - h * v
       return [vector, krylov_vectors, n, H]
 
+    def iterative_classical_gram_schmidt(
+        vector: jax.ShapedArray,
+        krylov_vectors: jax.ShapedArray,
+        iterations: int = 2) -> Tuple[jax.ShapedArray, jax.ShapedArray]:
+      """
+      Orthogonalize `vector`  to all rows of `krylov_vectors`, using
+      an iterated classical gram schmidt orthogonalization.
+      Args:
+        vector: Initial vector.
+        krylov_vectors: Matrix of krylov vectors, each row is treated as a
+          vector.
+        iterations: Number of iterations.
+      Returns:
+        jax.ShapedArray: The orthogonalized vector.
+        jax.ShapedArray: The overlaps of `vector` with all previous
+          krylov vectors
+      """
+      vec = vector
+      overlaps = 0
+      for _ in range(iterations):
+        ov = jax.numpy.dot(
+            krylov_vectors.conj(), vec, precision=precision)
+        vec = vec - jax.numpy.dot(
+            ov, krylov_vectors, precision=precision)
+        overlaps = overlaps + ov
+      return vec, overlaps
+
+    shape = v0.shape
+    Z = jax.numpy.linalg.norm(v0)
+    #only normalize if norm > eps, else return zero vector
+    v = jax.lax.cond(Z > eps, lambda x: v0 / Z, lambda x: v0 * 0.0, None)
+    krylov_vectors = krylov_vectors.at[start, :].set(jax.numpy.ravel(v))
+    H = jax.lax.cond(
+        start > 0,
+        lambda x: jax.ops.index_update(H, jax.ops.index[x, x - 1], Z),
+        lambda x: H, start)
     # body of the arnoldi iteration
     def body(vals):
-      krylov_vectors, H, matvec, vector, _, threshold, i, maxiter = vals
-      Av = matvec(vector, *args)
-      initial_vals = [Av, krylov_vectors, i, H]
-      Av, krylov_vectors, _, H = jax.lax.fori_loop(
-          0, i + 1, modified_gram_schmidt_step_arnoldi, initial_vals)
+      krylov_vectors, H, previous_vector, _, i = vals
+      Av = matvec(previous_vector, *args)
+
+      Av, overlaps = iterative_classical_gram_schmidt(
+          Av.ravel(),
+          (i >= jax.numpy.arange(krylov_vectors.shape[0]))[:, None] *
+          krylov_vectors)
+      H = H.at[:, i].set(overlaps)
       norm = jax.numpy.linalg.norm(Av)
-      Av /= norm
-      H = jax.ops.index_update(H, jax.ops.index[i + 1, i], norm)
-      krylov_vectors = jax.ops.index_update(krylov_vectors,
-                                            jax.ops.index[i + 1, :],
-                                            jax.numpy.ravel(Av))
-      return [krylov_vectors, H, matvec, Av, norm, threshold, i + 1, maxiter]
+      Av = jax.numpy.reshape(Av, shape)
+
+      # only normalize if norm is larger than threshold,
+      # otherwise return zero vector
+      Av = jax.lax.cond(norm > eps, lambda x: Av/norm, lambda x: Av * 0.0, None)
+      krylov_vectors, H = jax.lax.cond(
+          i < num_krylov_vecs - 1,
+          lambda x: (krylov_vectors.at[i + 1, :].set(Av.ravel()), H.at[i + 1, i].set(norm)), #pylint: disable=line-too-long
+          lambda x: (x[0], x[1]),
+          (krylov_vectors, H, Av, i, norm))
+
+      return [krylov_vectors, H, Av, norm, i + 1]
 
     def cond_fun(vals):
       # Continue loop while iteration < num_krylov_vecs and norm > eps
-      _, _, _, _, norm, _, iteration, _ = vals
+      norm, iteration = vals[3], vals[4]
       counter_done = (iteration >= num_krylov_vecs)
       norm_not_too_small = norm > eps
-      continue_iteration = jax.lax.cond(counter_done,
-                                        _, lambda x: False,
-                                        _, lambda x: norm_not_too_small)
-
+      continue_iteration = jax.lax.cond(counter_done, lambda x: False,
+                                        lambda x: norm_not_too_small, None)
       return continue_iteration
-    initial_norm = v.real.dtype.type(1.0+eps)
-    initial_values = [krylov_vectors, H, matvec, v, initial_norm, eps, start,
-                      num_krylov_vecs]
+
+    initial_values = [krylov_vectors, H, v, Z, start]
     final_values = jax.lax.while_loop(cond_fun, body, initial_values)
-    kvfinal, Hfinal, _, _, norm, _, it, _ = final_values
-    return kvfinal, Hfinal, it, norm < eps
+    krylov_vectors, H, residual, norm, it = final_values
+    return krylov_vectors, H, residual, norm, it, norm < eps
 
   return _arnoldi_fact
 
-
 def _implicitly_restarted_arnoldi(jax: types.ModuleType) -> Callable:
   """
   Helper function to generate a jitted function to do an
@@ -381,8 +431,7 @@ def _implicitly_restarted_arnoldi(jax: types.ModuleType) -> Callable:
       initial_state: An starting vector for the iteration.
       num_krylov_vecs: Number of krylov vectors of the arnoldi factorization.
         numeig: The number of desired eigenvector-eigenvalue pairs.
-      which: Which eigenvalues to target. Currently supported: `which = 'LR'`
-        or `which = 'LM'`.
+      which: Which eigenvalues to target. Currently supported: `which = 'LR'.
       eps: Convergence flag. If the norm of a krylov vector drops below `eps`
         the iteration is terminated.
       maxiter: Maximum number of (outer) iteration steps.
@@ -395,89 +444,90 @@ def _implicitly_restarted_arnoldi(jax: types.ModuleType) -> Callable:
     Callable: A function performing an implicitly restarted
       Arnoldi factorization
   """
+  JaxPrecisionType = type(jax.lax.Precision.DEFAULT)
 
   arnoldi_fact = _generate_arnoldi_factorization(jax)
 
   # ######################################################
   # #######  NEW SORTING FUCTIONS INSERTED HERE  #########
   # ######################################################
-  @functools.partial(jax.jit, static_argnums=(1,))
-  def LR_sort(evals, p):
-    inds = np.argsort(jax.numpy.real(evals), kind='stable')[::-1]
-    shifts = evals[inds][-p:]
-    return shifts, inds
-
-  @functools.partial(jax.jit, static_argnums=(1,))
-  def LM_sort(evals, p):
-    inds = np.argsort(jax.numpy.abs(evals), kind='stable')[::-1]
+  @functools.partial(jax.jit, static_argnums=(0,))
+  def LR_sort(
+      p: int,
+      evals: jax.ShapedArray) -> Tuple[jax.ShapedArray, jax.ShapedArray]:
+    inds = jax.numpy.argsort(jax.numpy.real(evals), kind='stable')[::-1]
     shifts = evals[inds][-p:]
     return shifts, inds
 
   # #######################################################
   # #######################################################
   # #######################################################
-  @functools.partial(jax.jit, static_argnums=(4, 5, 6))
-  def shifted_QR(Vm, Hm, fm, evals, k, p, which, res_thresh):
-    funs = [LR_sort, LM_sort]
-    shifts, _ = funs[which](evals, p)
-    # compress to k = numeig
-    q = jax.numpy.zeros(Hm.shape[0])
-    q = jax.ops.index_update(q, jax.ops.index[-1], 1)
-    m = Hm.shape[0]
-
-    for shift in shifts:
-      tmp = shift * jax.numpy.eye(m)
-      Qj, R = jax.numpy.linalg.qr(Hm - tmp)
-      Hm = R @ Qj + tmp
+  @functools.partial(jax.jit, static_argnums=(3, 4))
+  def shifted_QR(
+      Vm: jax.ShapedArray, Hm: jax.ShapedArray, fm: jax.ShapedArray,
+      numeig: int, sort_fun: Callable
+  ) -> Tuple[jax.ShapedArray, jax.ShapedArray, jax.ShapedArray]:
+    ######################################################
+    #######  NEW SORTING FUCTIONS INSERTED HERE  #########
+    ######################################################
+    ######################################################
+    ######################################################
+    evals, _ = jax.numpy.linalg.eig(Hm)
+    shifts, _ = sort_fun(evals)
+    # compress arnoldi factorization
+    q = jax.numpy.zeros(Hm.shape[0], dtype=Hm.dtype)
+    q = q.at[-1].set(1.0)
+
+    def body(i, vals):
+      Vm, Hm, q = vals
+      shift = shifts[i] * jax.numpy.eye(Hm.shape[0], dtype=Hm.dtype)
+      Qj, R = jax.numpy.linalg.qr(Hm - shift)
+      Hm = R @ Qj + shift
       Vm = Qj.T @ Vm
       q = q @ Qj
+      return Vm, Hm, q
 
-    fk = Vm[k, :] * Hm[k, k - 1] + fm * q[k - 1]
-    Vk = Vm[0:k, :]
-    Hk = Hm[0:k, 0:k]
-    H = jax.numpy.zeros((k + p + 1, k + p), dtype=fm.dtype)
-    H = jax.ops.index_update(H, jax.ops.index[0:k, 0:k], Hk)
-    Z = jax.numpy.linalg.norm(fk)
-    v = fk / Z
-    krylov_vectors = jax.numpy.zeros((k + p + 1, Vm.shape[1]), dtype=fm.dtype)
-    krylov_vectors = jax.ops.index_update(krylov_vectors, jax.ops.index[0:k, :],
-                                          Vk)
-    krylov_vectors = jax.ops.index_update(krylov_vectors, jax.ops.index[k:], v)
-    Z = jax.numpy.linalg.norm(fk)
-    #if fk is a zero-vector then arnoldi has exactly converged.
-    #use small threshold to check this
-    return krylov_vectors, H, fk, Z < res_thresh
-
-  @functools.partial(jax.jit, static_argnums=(2,))
-  def update_data(Vm_tmp, Hm_tmp, numits):
-    Vm = Vm_tmp[0:numits, :]
-    Hm = Hm_tmp[0:numits, 0:numits]
-    fm = Vm_tmp[numits, :] * Hm_tmp[numits, numits - 1]
-    return Vm, Hm, fm
+    Vm, Hm, q = jax.lax.fori_loop(0, shifts.shape[0], body,
+                                  (Vm, Hm, q))
+    fk = Vm[numeig, :] * Hm[numeig, numeig - 1] + fm * q[numeig - 1]
+    return Vm, Hm, fk
 
   @functools.partial(jax.jit, static_argnums=(3,))
-  def get_vectors(Vm, unitary, inds, numeig):
+  def get_vectors(Vm: jax.ShapedArray, unitary: jax.ShapedArray,
+                  inds: jax.ShapedArray, numeig: int) -> jax.ShapedArray:
 
-    def body_vector(i, vals):
-      krv, unitary, states, inds = vals
+    def body_vector(i, states):
       dim = unitary.shape[1]
       n, m = jax.numpy.divmod(i, dim)
       states = jax.ops.index_add(states, jax.ops.index[n, :],
-                                 krv[m, :] * unitary[m, inds[n]])
-      return [krv, unitary, states, inds]
+                                 Vm[m, :] * unitary[m, inds[n]])
+      return states
 
     state_vectors = jax.numpy.zeros([numeig, Vm.shape[1]], dtype=Vm.dtype)
-    _, _, state_vectors, _ = jax.lax.fori_loop(
-        0, numeig * Vm.shape[0], body_vector,
-        [Vm, unitary, state_vectors, inds])
+    state_vectors = jax.lax.fori_loop(0, numeig * Vm.shape[0], body_vector,
+                                      state_vectors)
     state_norms = jax.numpy.linalg.norm(state_vectors, axis=1)
     state_vectors = state_vectors / state_norms[:, None]
     return state_vectors
 
-
+  @functools.partial(jax.jit, static_argnums=(2, 3))
+  def check_eigvals_convergence_iram(beta_m: float, Hm: jax.ShapedArray,
+                                     eps: float, numeig: int) -> bool:
+    eigvals, eigvecs = jax.numpy.linalg.eig(Hm)
+    # TODO (mganahl) confirm that this is a valid matrix norm)
+    Hm_norm = jax.numpy.linalg.norm(Hm)
+    thresh = jax.numpy.maximum(
+        jax.numpy.finfo(eigvals.dtype).eps * Hm_norm,
+        jax.numpy.abs(eigvals[:numeig]) * eps)
+    vals = jax.numpy.abs(eigvecs[numeig - 1, :numeig])
+    return jax.numpy.all(beta_m * vals < thresh)
+
+  @functools.partial(jax.jit, static_argnums=(3, 4, 5, 6, 7, 8))
   def implicitly_restarted_arnoldi_method(
-      matvec, args, initial_state, num_krylov_vecs, numeig, which, eps, maxiter,
-      res_thresh, precision) -> Tuple[List[Tensor], List[Tensor]]:
+      matvec: Callable, args: List, initial_state: jax.ShapedArray,
+      num_krylov_vecs: int, numeig: int, which: Text, eps: float, maxiter: int,
+      precision: JaxPrecisionType
+  ) -> Tuple[jax.ShapedArray, List[jax.ShapedArray], int]:
     """
     Implicitly restarted arnoldi factorization of `matvec`. The routine
     finds the lowest `numeig` eigenvector-eigenvalue pairs of `matvec`
@@ -488,54 +538,67 @@ def implicitly_restarted_arnoldi_method(
     of `matvec` matches the dtype of the initial state. Otherwise jax
     will raise a TypeError.
 
+    NOTE: Under certain circumstances, the routine can return spurious 
+    eigenvalues 0.0: if the Arnoldi iteration terminated early
+    (after numits < num_krylov_vecs iterations)
+    and numeig > numits, then spurious 0.0 eigenvalues will be returned.
+
     Args:
       matvec: A callable representing the linear operator.
       args: Arguments to `matvec`.  `matvec` is called with
         `matvec(x, *args)` with `x` the input array on which
         `matvec` should act.
       initial_state: An starting vector for the iteration.
+      dim: The matrix dimension of the linear operator `matvec`.
       num_krylov_vecs: Number of krylov vectors of the arnoldi factorization.
         numeig: The number of desired eigenvector-eigenvalue pairs.
-      which: Which eigenvalues to target. Currently supported: `which = 'LR'`
-        or `which = 'LM'`.
+      which: Which eigenvalues to target. 
+        Currently supported: `which = 'LR'` (largest real part).
       eps: Convergence flag. If the norm of a krylov vector drops below `eps`
         the iteration is terminated.
       maxiter: Maximum number of (outer) iteration steps.
+      precision: jax.lax.Precision used within lax operations.
     Returns:
-      eta, U: Two lists containing eigenvalues and eigenvectors.
+      jax.ShapedArray: Eigenvalues
+      List: Eigenvectors
+      int: Number of inner krylov iterations of the last arnoldi 
+        factorization.
     """
-    N = np.prod(initial_state.shape)
-    p = num_krylov_vecs - numeig
-    num_krylov_vecs = np.min([num_krylov_vecs, N])
-    if (p <= 1) and (num_krylov_vecs < N):
-      raise ValueError(f"`num_krylov_vecs` must be between `numeig` + 1 <"
-                       f" `num_krylov_vecs` <= N={N},"
-                       f" `num_krylov_vecs`={num_krylov_vecs}")
-
+    shape = initial_state.shape
     dtype = initial_state.dtype
+
+    dim = np.prod(shape).astype(np.int32)
+    num_expand = num_krylov_vecs - numeig
+
+    if (num_expand <= 1) and (num_krylov_vecs < dim):
+      raise ValueError(f"num_krylov_vecs must be between numeig + 1 <"
+                       f" num_krylov_vecs <= dim = {dim},"
+                       f" num_krylov_vecs = {num_krylov_vecs}")
+    if numeig > dim:
+      raise ValueError(f"number of requested eigenvalues numeig = {numeig} "
+                       f"is larger than the dimension of the operator "
+                       f"dim = {dim}")
+
     # initialize arrays
-    krylov_vectors = jax.numpy.zeros(
-        (num_krylov_vecs + 1, jax.numpy.ravel(initial_state).shape[0]),
-        dtype=dtype)
-    H = jax.numpy.zeros((num_krylov_vecs + 1, num_krylov_vecs), dtype=dtype)
+    Vm = jax.numpy.zeros(
+        (num_krylov_vecs, jax.numpy.ravel(initial_state).shape[0]), dtype=dtype)
+    Hm = jax.numpy.zeros((num_krylov_vecs, num_krylov_vecs), dtype=dtype)
     # perform initial arnoldi factorization
-    Vm_tmp, Hm_tmp, numits, converged = arnoldi_fact(matvec, args,
-                                                     initial_state,
-                                                     krylov_vectors, H, 0,
-                                                     num_krylov_vecs, eps,
-                                                     precision)
-    # obtain an m-step arnoldi factorization
-    Vm, Hm, fm = update_data(Vm_tmp, Hm_tmp, numits)
-
-    it = 0
+    Vm, Hm, residual, norm, numits, ar_converged = arnoldi_fact(
+        matvec, args, initial_state, Vm, Hm, 0, num_krylov_vecs, eps, precision)
+    fm = residual.ravel() * norm
+
+
+    # sort_fun returns `num_expand` least relevant eigenvalues
+    # (those to be projected out)
     if which == 'LR':
-      _which = 0
-    elif which == 'LM':
-      _which = 1
+      sort_fun = jax.tree_util.Partial(functools.partial(LR_sort, num_expand))
     else:
       raise ValueError(f"which = {which} not implemented")
-    # make sure the dtypes are matching
-    if maxiter > 0:
+
+    it = 1  # we already did one arnoldi factorization
+    if maxiter > 1:
+      # cast arrays to correct complex dtype
       if Vm.dtype == np.float64:
         dtype = np.complex128
       elif Vm.dtype == np.float32:
@@ -546,33 +609,89 @@ def implicitly_restarted_arnoldi_method(
         dtype = Vm.dtype
       else:
         raise TypeError(f'dtype {Vm.dtype} not supported')
+
       Vm = Vm.astype(dtype)
       Hm = Hm.astype(dtype)
       fm = fm.astype(dtype)
 
-    while (it < maxiter) and (not converged):
-      evals, _ = jax.numpy.linalg.eig(Hm)
-      krylov_vectors, H, fk, converged = shifted_QR(Vm, Hm, fm, evals, numeig,
-                                                    p, _which, res_thresh)
-      if converged:
-        break
-      v0 = jax.numpy.reshape(fk, initial_state.shape)
-      # restart
-      Vm_tmp, Hm_tmp, _, converged = arnoldi_fact(matvec, args, v0,
-                                                  krylov_vectors, H, numeig,
-                                                  num_krylov_vecs, eps,
-                                                  precision)
-      Vm, Hm, fm = update_data(Vm_tmp, Hm_tmp, num_krylov_vecs)
-      it += 1
-
+    def outer_loop(carry):
+      Hm, Vm, fm, it, numits, ar_converged, _, _, = carry
+      # perform shifted QR iterations to compress arnoldi factorization
+      # Note that ||fk|| typically decreases as one iterates the outer loop
+      # indicating that iram converges.
+      # ||fk|| = \beta_m in reference above
+      Vk, Hk, fk = shifted_QR(Vm, Hm, fm, numeig, sort_fun)
+      # reset matrices
+      Vk = Vk.at[numeig:, :].set(0.0)
+      Hk = Hk.at[numeig:, :].set(0.0)
+      Hk = Hk.at[:, numeig:].set(0.0)
+      beta_k = jax.numpy.linalg.norm(fk)
+      converged = check_eigvals_convergence_iram(beta_k, Hk, eps, numeig)
+
+      def do_arnoldi(vals):
+        Vk, Hk, fk, _, _, _, _ = vals
+        # restart
+        Vm, Hm, residual, norm, numits, ar_converged = arnoldi_fact(
+            matvec, args, jax.numpy.reshape(fk, shape), Vk, Hk, numeig,
+            num_krylov_vecs, eps, precision)
+        fm = residual.ravel() * norm
+        return [Vm, Hm, fm, norm, numits, ar_converged, False]
+
+      def cond_arnoldi(vals):
+        return vals[6]
+
+      res = jax.lax.while_loop(cond_arnoldi, do_arnoldi, [
+          Vk, Hk, fk,
+          jax.numpy.linalg.norm(fk), numeig, False,
+          jax.numpy.logical_not(converged)
+      ])
+
+
+
+      Vm, Hm, fm, norm, numits, ar_converged = res[0:6]
+      out_vars = [
+          Hm, Vm, fm, it + 1, numits, ar_converged, converged, norm
+      ]
+      return out_vars
+
+    def cond_fun(carry):
+      it, ar_converged, converged = carry[3], carry[5], carry[
+          6]
+      return jax.lax.cond(
+          it < maxiter, lambda x: x, lambda x: False,
+          jax.numpy.logical_not(jax.numpy.logical_or(converged, ar_converged)))
+
+    converged = False
+    carry = [Hm, Vm, fm, it, numits, ar_converged, converged, norm]
+    res = jax.lax.while_loop(cond_fun, outer_loop, carry)
+    Hm, Vm = res[0], res[1]
+    numits, converged = res[4], res[6]
+    # if `ar_converged` then `norm`is below convergence threshold
+    # set it to 0.0 in this case to prevent `jnp.linalg.eig` from finding a
+    # spurious eigenvalue of order `norm`.
+    Hm = Hm.at[numits, numits - 1].set(
+        jax.lax.cond(converged, lambda x: Hm.dtype.type(0.0), lambda x: x,
+                     Hm[numits, numits - 1]))
+
+    # if the Arnoldi-factorization stopped early (after `numit` iterations)
+    # before exhausting the allowed size of the Krylov subspace,
+    # (i.e. `numit` < 'num_krylov_vecs'), set elements
+    # at positions m, n with m, n >= `numit` to 0.0.
+
+    # FIXME (mganahl): under certain circumstances, the routine can still
+    # return spurious 0 eigenvalues: if arnoldi terminated early
+    # (after numits < num_krylov_vecs iterations)
+    # and numeig > numits, then spurious 0.0 eigenvalues will be returned
+
+    Hm = (numits > jax.numpy.arange(num_krylov_vecs))[:, None] * Hm * (
+        numits > jax.numpy.arange(num_krylov_vecs))[None, :]
     eigvals, U = jax.numpy.linalg.eig(Hm)
-    _, inds = LR_sort(eigvals, _which)
+    inds = sort_fun(eigvals)[1][:numeig]
     vectors = get_vectors(Vm, U, inds, numeig)
-
-    return eigvals[inds[0:numeig]], [
-        jax.numpy.reshape(vectors[n, :], initial_state.shape)
+    return eigvals[inds], [
+        jax.numpy.reshape(vectors[n, :], shape)
         for n in range(numeig)
-    ]
+    ], numits
 
   return implicitly_restarted_arnoldi_method
 


Test patch:
diff --git a/tensornetwork/backends/jax/jax_backend_test.py b/tensornetwork/backends/jax/jax_backend_test.py
--- a/tensornetwork/backends/jax/jax_backend_test.py
+++ b/tensornetwork/backends/jax/jax_backend_test.py
@@ -569,7 +569,7 @@ def compare_eigvals_and_eigvecs(U, eta, U_exact, eta_exact, thresh=1E-8):
 
 
 @pytest.mark.parametrize("dtype", [np.float64, np.complex128])
-@pytest.mark.parametrize("which", ["LR", "LM"])
+@pytest.mark.parametrize("which", ["LR"])
 def test_eigs_all_eigvals_with_init(dtype, which):
   backend = jax_backend.JaxBackend()
   D = 16
@@ -587,7 +587,7 @@ def mv(x, H):
 
 
 @pytest.mark.parametrize("dtype", [np.float64, np.complex128])
-@pytest.mark.parametrize("which", ["LR", "LM"])
+@pytest.mark.parametrize("which", ["LR"])
 def test_eigs_all_eigvals_no_init(dtype, which):
   backend = jax_backend.JaxBackend()
   D = 16
@@ -610,7 +610,7 @@ def mv(x, H):
 
 
 @pytest.mark.parametrize("dtype", [np.float64, np.complex128])
-@pytest.mark.parametrize("which", ["LR", "LM"])
+@pytest.mark.parametrize("which", ["LR"])
 def test_eigs_few_eigvals_with_init(dtype, which):
   backend = jax_backend.JaxBackend()
   D = 16
@@ -629,7 +629,7 @@ def mv(x, H):
 
 
 @pytest.mark.parametrize("dtype", [np.float64, np.complex128])
-@pytest.mark.parametrize("which", ["LR", "LM"])
+@pytest.mark.parametrize("which", ["LR"])
 def test_eigs_few_eigvals_no_init(dtype, which):
   backend = jax_backend.JaxBackend()
   D = 16
@@ -652,7 +652,7 @@ def mv(x, H):
 
 
 @pytest.mark.parametrize("dtype", [np.float64, np.complex128])
-@pytest.mark.parametrize("which", ["LR", "LM"])
+@pytest.mark.parametrize("which", ["LR"])
 def test_eigs_large_ncv_with_init(dtype, which):
   backend = jax_backend.JaxBackend()
   D = 16
@@ -699,7 +699,7 @@ def test_eigs_raises():
         ValueError, match=f"which = {which}"
         f" is currently not supported."):
       backend.eigs(lambda x: x, which=which)
-  with pytest.raises(KeyError, match="dtype"):
+  with pytest.raises(TypeError, match="dtype"):
     backend.eigs(lambda x: x, shape=(10,), dtype=np.int32)
 
 ##################################################################
diff --git a/tensornetwork/backends/jax/jitted_functions_test.py b/tensornetwork/backends/jax/jitted_functions_test.py
--- a/tensornetwork/backends/jax/jitted_functions_test.py
+++ b/tensornetwork/backends/jax/jitted_functions_test.py
@@ -7,6 +7,7 @@
 jax_dtypes = [np.float32, np.float64, np.complex64, np.complex128]
 precision = jax.lax.Precision.HIGHEST
 
+
 @pytest.mark.parametrize("dtype", [np.float64, np.complex128])
 def test_arnoldi_factorization(dtype):
   np.random.seed(10)
@@ -21,18 +22,18 @@ def matvec(vector, matrix):
 
   arnoldi = jitted_functions._generate_arnoldi_factorization(jax)
   ncv = 40
-  kv = jax.numpy.zeros((ncv + 1, D), dtype=dtype)
-  H = jax.numpy.zeros((ncv + 1, ncv), dtype=dtype)
+  kv = jax.numpy.zeros((ncv, D), dtype=dtype)
+  H = jax.numpy.zeros((ncv, ncv), dtype=dtype)
   start = 0
-  tol = 0.01
-  kv, H, it, _ = arnoldi(matvec, [mat], x, kv, H, start, ncv, tol, precision)
-  Vm = jax.numpy.transpose(kv[:it, :])
-  Hm = H[:it, :it]
-  fm = kv[it, :] * H[it, it - 1]
-  em = np.zeros((1, Vm.shape[1]))
+  tol = 1E-5
+  Vm, Hm, residual, norm, it, _ = arnoldi(matvec, [mat], x, kv, H, start, ncv,
+                                          tol, precision)
+  fm = residual * norm
+  em = np.zeros((1, Vm.shape[0]))
   em[0, -1] = 1
-  np.testing.assert_almost_equal(mat @ Vm - Vm @ Hm - fm[:, None] * em,
-                                 np.zeros((it, Vm.shape[1])).astype(dtype))
+  #test arnoldi relation
+  np.testing.assert_almost_equal(mat @ Vm.T - Vm.T @ Hm - fm[:, None] * em,
+                                 np.zeros((it, Vm.shape[0])).astype(dtype))
 
 
 @pytest.mark.parametrize("dtype", jax_dtypes)