Jax gmres by alewis · Pull Request #707 · google/TensorNetwork

alewis · 2020-07-02T21:06:53Z

Adds a Jax implementation of GMRES.

mganahl · 2020-07-02T21:08:49Z

+  return x
+
+
+@partial(jax.jit, static_argnums=(2,))


Is this returning an arnoldi factorization?

If so, do you think you can use the one that is implemented in the backend?

Well, the one in the backend returns the eigendecomposition; I need the V and H from the Krylov decomposition (actually I need V and the full QR decomposition of H). If Arnoldi were refactored to produce these it could be used.

However, a more efficient implementation, which I will add later, would also update the GMRES residual by a Givens rotation once per new Krylov vector. This becomes specialized enough that I think it makes more sense to keep things together with GMRES until proved otherwise, which is reflected in the code as written.

sorry I was not precise in my wording. I meant the arnoldi factorization in jitted_functions.py. This one priduces exactly that. It' called _arnoldi_fact

would it make sense to move gmres into the jitted_functions.py file?

My recent push should address the above

chaserileyroberts

LGTM readability and design wise. Deferring to Martin's judgement for correctness.

alewis · 2020-07-07T18:05:10Z

The build works now

mganahl · 2020-07-07T18:16:30Z

-      self._iram = jitted_functions._implicitly_restarted_arnoldi(libjax)
-    return self._iram(_CACHED_MATVECS[A], args, initial_state, num_krylov_vecs,
-                      numeig, which, tol, maxiter)
+    imp_arnoldi = jitted_functions._implicitly_restarted_arnoldi(libjax)


love that name!

mganahl · 2020-07-07T18:18:58Z

  """
-  Helper function to generate jitted lanczos function used 
-  in JaxBackend.eigsh_lanczos. The function `jax_lanczos` 
+  Helper function to generate jitted lanczos function used


Did you remove trailing white space on purpose, or is this from yapfing?

I removed it on purpose with a vi script

Trailing whitespace violates PEP8 and is in general Bad because for example it invisible modifies the behaviour of text editor commands to jump to the end of the line. Was there some reason to keep it?

no, just curious. If you want it removed, then it would be great if you could prepare a PR that fixes this for the whole library (just run your script over all files). This way we'll have a cleaner code history

mganahl · 2020-07-07T18:36:41Z

+                                        _, lambda x: norm_not_too_small)
+
+      return continue_iteration
+    initial_norm_typecaster = np.zeros((1,)) + eps + 1.


you can do initial_norm = v.real.dtype.type(1.0+eps) instead

mganahl · 2020-07-07T18:38:00Z

+    Given a linear mapping with (n x n) matrix representation
+        A = A_mv(*A_args) gmres_m solves
+        Ax = b          (1)
+    where x and b are length-b vectors, using the method of


length-n vectors?

mganahl · 2020-07-07T18:38:37Z

+
+    A_mv     : A function `v0 = A_mv(v, *A_args, **A_kwargs)` where `v0` and
+               `v` have the same shape.
+    b        : The `b` in `A @ x = b`.


reverse order b and A_args

looks like A_args is missing from the docstring

mganahl · 2020-07-07T18:40:18Z

+    return r, beta
+
+
+  #@partial(jax.jit, static_argnums=(2,))


why is the jit command commented?

mganahl · 2020-07-07T19:18:51Z

+    backend.gmres(dummy_mv, b, A_kwargs=A_kwargs)
+
+
+#jax_qr_dtypes = [np.float32, np.float64, np.complex64, np.complex128]


remove commented code

mganahl · 2020-07-07T19:19:32Z

+
+
+#jax_qr_dtypes = [np.float32, np.float64, np.complex64, np.complex128]
+jax_qr_dtypes = [np.float32]


Also test for float64

and complex if possible

mganahl · 2020-07-07T19:20:13Z

+  x, _ = backend.gmres(A_mv, b, x0=x0, num_krylov_vectors=n_kry, tol=tol)
+  solution = jax.numpy.array([2., 1.], dtype=dtype)
+  eps = jax.numpy.linalg.norm(jax.numpy.abs(solution) - jax.numpy.abs(x))
+  print(eps)


remove print

mganahl · 2020-07-07T19:21:39Z

+  def A_mv(x):
+    return A @ x
+  b = A_mv(solution)
+  tol = b.size * jax.numpy.finfo(dtype).eps


good to know!

alewis · 2020-07-07T20:03:39Z

Martin's comments should be dealt with

mganahl

Can you add the A_args to the docstring, then this is good to go

mganahl · 2020-07-08T14:03:19Z

Hey @alewis, did you check that your modifications to eigsh_lanczos and eigs do not trigger unneccesary tracing? I checkout out your PR and on my laptop it seems it does

alewis · 2020-07-08T15:32:06Z

I find that tracing is triggered by the first run only, as in the screenshot. However, running the cell twice does trace twice; this also appears to have been true in the older code.

mganahl · 2020-07-08T15:39:43Z

import jax
import tensornetwork as tn
import numpy as np
import time
D=200
dtype=np.float32
matrix = jax.numpy.array(np.random.rand(D,D).astype(dtype))
vector = jax.numpy.array(np.random.rand(D,).astype(dtype))
@jax.jit
def matvec_jax_matrix(vec,matrix):
    return jax.numpy.tensordot(matrix, vec,([1],[0]))
jax_backend = tn.backends.jax.jax_backend.JaxBackend()
ncv=10
t1 = time.time()
eta_j, U_j = jax_backend.eigsh_lanczos(matvec_jax_matrix,[matrix],vector,num_krylov_vecs = ncv,numeig=1, 
                                       reorthogonalize=False)
print('jax eigvals:', eta_j)
t2 = time.time()
eta_j, U_j = jax_backend.eigsh_lanczos(matvec_jax_matrix,[matrix],vector,num_krylov_vecs = ncv,numeig=1, 
                                       reorthogonalize=False)

print('jax eigvals:', eta_j)
t3 = time.time()
print('jax first:', t2 - t1)
print('jax second:', t3 - t2)

could you run this and send the result?

alewis · 2020-07-08T15:45:05Z

[image: Screen Shot 2020-07-08 at 11.44.19 AM.png]

…

On Wed, Jul 8, 2020 at 11:39 AM Martin Ganahl ***@***.***> wrote: import jaximport tensornetwork as tnimport numpy as npimport timeD=200dtype=np.float32matrix = jax.numpy.array(np.random.rand(D,D).astype(dtype))vector = ***@***.*** matvec_jax_matrix(vec,matrix): return jax.numpy.tensordot(matrix, vec,([1],[0]))jax_backend = tn.backends.jax.jax_backend.JaxBackend()ncv=10t1 = time.time()eta_j, U_j = jax_backend.eigsh_lanczos(matvec_jax_matrix,[matrix],vector,num_krylov_vecs = ncv,numeig=1, reorthogonalize=False)print('jax eigvals:', eta_j)t2 = time.time()eta_j, U_j = jax_backend.eigsh_lanczos(matvec_jax_matrix,[matrix],vector,num_krylov_vecs = ncv,numeig=1, reorthogonalize=False) print('jax eigvals:', eta_j)t3 = time.time()print('jax first:', t2 - t1)print('jax second:', t3 - t2) could you run this and send the result? — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#707 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAINHBQ64NFTMLIT64DEJOLR2SHM5ANCNFSM4OPHYVLQ> .

mganahl · 2020-07-08T15:46:20Z

sorry can't see it

mganahl · 2020-07-08T16:34:06Z

looks good, I'll pull it in as soon as building's finished

alewis added 6 commits June 1, 2020 12:27

Add linalg directory and initialization code

81b9db3

Delete linalg for merge

c52b16b

Merge https://github.com/google/TensorNetwork

15f9447

Merge remote-tracking branch 'upstream/master'

f826978

Merge remote-tracking branch 'upstream/master'

d547918

Jax GMRES and tests

662ac2f

googlebot added the cla: yes label Jul 2, 2020

mganahl reviewed Jul 2, 2020

View reviewed changes

alewis added 3 commits July 3, 2020 11:15

Correct typing of ShapedArray

bbb7849

Correct typing in gmres.py

dc47792

Refactor GMRES to use existing Arnoldi

e6136cd

chaserileyroberts approved these changes Jul 6, 2020

View reviewed changes

Merge branch 'master' into jax_gmres

6c3bb62

kshithijiyer approved these changes Jul 7, 2020

View reviewed changes

Fix obsolete test from merge

07eecfb

mganahl reviewed Jul 7, 2020

View reviewed changes

Comment thread tensornetwork/backends/jax/jax_backend.py

mganahl reviewed Jul 7, 2020

View reviewed changes

Comment thread tensornetwork/backends/jax/jax_backend_test.py

mganahl reviewed Jul 7, 2020

View reviewed changes

Cleanups for PR

e0b85b9

mganahl approved these changes Jul 7, 2020

View reviewed changes

Correct docstring

8d117cd

Fix caching

dda87d0

chaserileyroberts merged commit db81bc0 into google:master Jul 8, 2020

		backend.gmres(dummy_mv, b, A_kwargs=A_kwargs)


		#jax_qr_dtypes = [np.float32, np.float64, np.complex64, np.complex128]



		#jax_qr_dtypes = [np.float32, np.float64, np.complex64, np.complex128]
		jax_qr_dtypes = [np.float32]

Uh oh!

Conversation

alewis commented Jul 2, 2020

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mganahl Jul 3, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chaserileyroberts left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alewis commented Jul 7, 2020

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alewis commented Jul 7, 2020

Uh oh!

mganahl left a comment

Choose a reason for hiding this comment

Uh oh!

mganahl commented Jul 8, 2020

Uh oh!

alewis commented Jul 8, 2020

Uh oh!

mganahl commented Jul 8, 2020

Uh oh!

alewis commented Jul 8, 2020 via email

Uh oh!

mganahl commented Jul 8, 2020

Uh oh!

mganahl commented Jul 8, 2020

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

mganahl Jul 3, 2020 •

edited

Loading

chaserileyroberts left a comment •

edited

Loading