fix error check in jax eigs (#909)

* fix error check in jax eigs

* fix docstring of krylov.eigs and eigsh_lanczos

* fix bug, update docstrings

* fix tests

* fix test

Author: mganahl

tensornetwork/backends/jax/jax_backend.py

@@ -281,7 +281,7 @@ def A(H,x):
       A: A (sparse) implementation of a linear operator.
          Call signature of `A` is `res = A(vector, *args)`, where `vector`
          can be an arbitrary `Tensor`, and `res.shape` has to be `vector.shape`.
-      arsg: A list of arguments to `A`.  `A` will be called as
+      args: A list of arguments to `A`.  `A` will be called as
         `res = A(initial_state, *args)`.
       initial_state: An initial vector for the algorithm. If `None`,
         a random initial `Tensor` is created using the `backend.randn` method

tensornetwork/backends/jax/jax_backend_test.py

@@ -748,7 +748,7 @@ def mv(x, H):
 def test_eigs_eigsh_large_ncv_with_init(dtype, solver, matrix_generator,
                                         exact_decomp, which):
   backend = jax_backend.JaxBackend()
-  D = 16
+  D = 100
   np.random.seed(10)
   init = backend.randn((D,), dtype=dtype, seed=10)
   H = matrix_generator(backend, dtype, D)
@@ -949,7 +949,8 @@ def test_eigs_eigsh_raises(solver, whichs):
 def test_eigs_dtype_raises():
   solver = jax_backend.JaxBackend().eigs
   with pytest.raises(TypeError, match="dtype"):
-    solver(lambda x: x, shape=(10,), dtype=np.int32)
+    solver(lambda x: x, shape=(10,), dtype=np.int32,
+           num_krylov_vecs=10)
 
 ##################################################################
 #############  This test should just not crash    ################

tensornetwork/backends/jax/jitted_functions.py

@@ -743,10 +743,11 @@ def implicitly_restarted_arnoldi_method(
 
     dim = np.prod(shape).astype(np.int32)
     num_expand = num_krylov_vecs - numeig
-    if num_krylov_vecs <= numeig < dim:
-      raise ValueError(f"num_krylov_vecs must be between numeig <"
-                       f" num_krylov_vecs <= dim = {dim},"
-                       f" num_krylov_vecs = {num_krylov_vecs}")
+    if not numeig <= num_krylov_vecs <= dim:
+      raise ValueError(f"num_krylov_vecs must be between numeig <="
+                       f" num_krylov_vecs <= dim, got "
+                       f" numeig = {numeig}, num_krylov_vecs = "
+                       f"{num_krylov_vecs}, dim = {dim}.")
     if numeig > dim:
       raise ValueError(f"number of requested eigenvalues numeig = {numeig} "
                        f"is larger than the dimension of the operator "

tensornetwork/linalg/krylov.py

@@ -125,11 +125,13 @@ def eigsh_lanczos(A: Callable,
   """
   Lanczos method for finding the lowest eigenvector-eigenvalue pairs
   of `A`.
+
   Args:
     A: A (sparse) implementation of a linear operator.
        Call signature of `A` is `res = A(vector, *args)`, where `vector`
        can be an arbitrary `Array`, and `res.shape` has to be `vector.shape`.
-    arsg: A list of arguments to `A`.  `A` will be called as
+    backend: A backend, text specifying one, or None.
+    args: A list of arguments to `A`.  `A` will be called as
       `res = A(x0, *args)`.
     x0: An initial vector for the Lanczos algorithm. If `None`,
       a random initial vector is created using the `backend.randn` method
@@ -152,6 +154,7 @@ def eigsh_lanczos(A: Callable,
       iterations to check convergence.
     reorthogonalize: If `True`, Krylov vectors are kept orthogonal by
       explicit orthogonalization (more costly than `reorthogonalize=False`)
+
   Returns:
     (eigvals, eigvecs)
      eigvals: A list of `numeig` lowest eigenvalues
@@ -182,39 +185,70 @@ def eigs(A: Callable,
          which: Text = 'LR',
          maxiter: int = 20) -> Tuple[Tensor, List]:
   """
-  Lanczos method for finding the lowest eigenvector-eigenvalue pairs
-  of `A`.
+  Implicitly restarted Arnoldi method for finding the lowest
+  eigenvector-eigenvalue pairs of a linear operator `A`.
+  `A` is a function implementing the matrix-vector
+  product.
+
+  WARNING: This routine uses jax.jit to reduce runtimes. jitting is triggered
+  at the first invocation of `eigs`, and on any subsequent calls
+  if the python `id` of `A` changes, even if the formal definition of `A`
+  stays the same.
+  Example: the following will jit once at the beginning, and then never again:
+
+  ```python
+  import jax
+  import numpy as np
+  def A(H,x):
+    return jax.np.dot(H,x)
+  for n in range(100):
+    H = jax.np.array(np.random.rand(10,10))
+    x = jax.np.array(np.random.rand(10,10))
+    res = eigs(A, [H],x) #jitting is triggerd only at `n=0`
+  ```
+
+  The following code triggers jitting at every iteration, which
+  results in considerably reduced performance
+
+  ```python
+  import jax
+  import numpy as np
+  for n in range(100):
+    def A(H,x):
+      return jax.np.dot(H,x)
+    H = jax.np.array(np.random.rand(10,10))
+    x = jax.np.array(np.random.rand(10,10))
+    res = eigs(A, [H],x) #jitting is triggerd at every step `n`
+  ```
+
   Args:
     A: A (sparse) implementation of a linear operator.
        Call signature of `A` is `res = A(vector, *args)`, where `vector`
-       can be an arbitrary `Array`, and `res.shape` has to be `vector.shape`.
-    arsg: A list of arguments to `A`.  `A` will be called as
-      `res = A(x0, *args)`.
-    x0: An initial vector for the Lanczos algorithm. If `None`,
-      a random initial vector is created using the `backend.randn` method
+       can be an arbitrary `Tensor`, and `res.shape` has to be `vector.shape`.
+    backend: A backend, text specifying one, or None.
+    args: A list of arguments to `A`.  `A` will be called as
+      `res = A(initial_state, *args)`.
+    x0: An initial vector for the algorithm. If `None`,
+      a random initial `Tensor` is created using the `backend.randn` method
     shape: The shape of the input-dimension of `A`.
-    dtype: The dtype of the input `A`. If both no `x0` is provided,
+    dtype: The dtype of the input `A`. If no `initial_state` is provided,
       a random initial state with shape `shape` and dtype `dtype` is created.
     num_krylov_vecs: The number of iterations (number of krylov vectors).
-    numeig: The nummber of eigenvector-eigenvalue pairs to be computed.
-      If `numeig > 1`, `reorthogonalize` has to be `True`.
-    tol: The desired precision of the eigenvalus. Uses
-      `backend.norm(eigvalsnew[0:numeig] - eigvalsold[0:numeig]) < tol`
-      as stopping criterion between two diagonalization steps of the
-      tridiagonal operator.
-    delta: Stopping criterion for Lanczos iteration.
-      If a Krylov vector :math: `x_n` has an L2 norm
-      :math:`\\lVert x_n\\rVert < delta`, the iteration
-      is stopped. It means that an (approximate) invariant subspace has
-      been found.
-    ndiag: The tridiagonal Operator is diagonalized every `ndiag`
-      iterations to check convergence.
-    reorthogonalize: If `True`, Krylov vectors are kept orthogonal by
-      explicit orthogonalization (more costly than `reorthogonalize=False`)
+    numeig: The number of eigenvector-eigenvalue pairs to be computed.
+    tol: The desired precision of the eigenvalues. For the jax backend
+      this has currently no effect, and precision of eigenvalues is not
+      guaranteed. This feature may be added at a later point. To increase
+      precision the caller can either increase `maxiter` or `num_krylov_vecs`.
+    which: Flag for targetting different types of eigenvalues. Currently
+      supported are `which = 'LR'` (larges real part) and `which = 'LM'`
+      (larges magnitude).
+    maxiter: Maximum number of restarts. For `maxiter=0` the routine becomes
+      equivalent to a simple Arnoldi method.
+
   Returns:
     (eigvals, eigvecs)
-     eigvals: A list of `numeig` lowest eigenvalues
-     eigvecs: A list of `numeig` lowest eigenvectors
+     eigvals: A list of `numeig` eigenvalues
+     eigvecs: A list of `numeig` eigenvectors
   """
   backend, x0_array, args_array = krylov_error_checks(backend, x0, args)
   mv = KRYLOV_MATVEC_CACHE.retrieve(backend.name, A)