Jax.sparse.linalg.eigsh like scipy eigsh?

[checc1]

Hello, I was wondering if a method such as scipy.sparse.linalg. eigsh exist in jax to find n eigenvalues and eigenvectors of a square Hermitian matrix.
I mean, a possibility of vectorizing in jax-like style this procedure on a sparse matrix to speed it up (I am aware that storing the entire matrix and placing it into a jax.array() would slow down the process due to memory/storing reasons).

[jakevdp]

You could try the experimental iterative eigenvalue solver here: https://docs.jax.dev/en/latest/_autosummary/jax.experimental.sparse.linalg.lobpcg_standard.html

[checc1]

Thank you for the quick reply. I am going to try it. 1 more question:
if I need to find the k eigenvectors within an iterative loop (Hamiltonian time evolution) where the matrix is evaluated at different subsequent time-steps, would it be possible to combine this eigenvalue solver with the jax.vmap(function, axes)) where the function takes in the time-step t[i] and the state psi , or I am off the track?

[jakevdp]

Yes, I believe lobpcg_standard is compatible with vmap