Computing the leading singular vector of a matrix without computing the full SVD - analog of scipy.sparse.linalg.svds

[tonysf]

I cannot find a function in jax that computes the leading singular vector of a matrix without also computing the full singular value decomposition. Essentially, I am looking for jax analog of the scipy function scipy.sparse.linalg.svds when k=1. Does such a function exist already in jax?

[f0uriest]

I think under the hood svds is just calling eighs on A.T@A, you could probably try something similar with jax.experimental.sparse.linalg.lobpcg_standard

[lobpcg]

Just follow the logic of scipy.sparse.linalg.svds with solver lobpcg