-
Notifications
You must be signed in to change notification settings - Fork 2.7k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Adds matrix sqrt #9544
Adds matrix sqrt #9544
Conversation
213d9a6
to
fbff65c
Compare
Thanks for putting together! Matrix square root is definitely very welcome functionality for JAX. How does performance compare to SciPy? I am a little concerned that iterating over individual matrix elements in JAX is going to be painfully slow. Usually we rely on compiled routines for this sort of thing, e.g., from Lapack or CuSolver. |
I ran a few tests (in double precision) and it performs as well as the non-blocked version of The execution times are the average of 5 runs of the following functions: def time_jax(n):
x = (np.random.randn(n,n))
tic = time.perf_counter()
jsp.linalg.sqrtm(x).block_until_ready()
return time.perf_counter() - tic
def time_scipy(n):
x = (np.random.randn(n,n))
tic = time.perf_counter()
sp.linalg.sqrtm(x, blocksize=1)
return time.perf_counter() - tic
def time_scipy_sylvester(n):
x = (np.random.randn(n,n))
tic = time.perf_counter()
sp.linalg.sqrtm(x)
return time.perf_counter() - tic |
I agree with Stephan that it is likely the scalar-loop formulation isn't ideal, especially on GPU or TPU where scalar code is very inefficient. But... we don't have a schur decomposition on GPU or TPU anyway. So that likely doesn't have a huge practical impact. If it were possible to avoid the scalar formulation, it would be preferable, but I don't think it should block merging this. |
I thought some more about this. It does feel un-jax like to have these nested impure loops and it does become slow for matrices of size exceeding This version is only on CPU and does not yet have a jvp defined (although the jvp is only a sylvester solver away). Because of these reasons, there might not be so much urgency to having it included in JAX. So a possible choice is to table this PR and instead work towards writing the custom call or a sylvester solver to make this |
jax/_src/scipy/linalg.py
Outdated
@@ -103,7 +103,23 @@ def eigh(a, b=None, lower=True, eigvals_only=False, overwrite_a=False, | |||
del overwrite_a, overwrite_b, turbo, check_finite | |||
return _eigh(a, b, lower, eigvals_only, eigvals, type) | |||
|
|||
@partial(jit, static_argnames=('output',)) | |||
def _schur(a, output): | |||
if output != "real": |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This should raise an error if an invalid output
is provided
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I added a line to check the output
argument schur
.
I think this implementation is actually probably OK for a first pass. Performance within a factor of 2-4x of SciPy is not terrible and I think matrices of size 5000x5000 are very rare. JVPs and even VJPs should work fine out of the box, because you wrote this in terms of existing JAX primitives, though it's possible there are more numerically stable implementations of gradients. |
c5a0c32
to
624236d
Compare
34ac130
to
7289ffe
Compare
7289ffe
to
cb73232
Compare
This PR adds an implementation of the matrix square root (makes progress on #2478). The algorithm is a Schur decomposition based method and is entirely performed with complex types. (Recent improvements of this method use Sylvester equation solvers which aren't available in JAX for the moment as mentioned in #6089)
To make things jittable, the main routine, computing the square root of upper-triangular matrices, uses
lax.fori_loop
which makes things a bit obscure. Before the jit-friendly-rewrite, the function here looked like this:(unsure if the proposed jit-friendly-rewrite is best)
I didn't squash the two commits for clarity. The first wraps the existing schur primitive into a
jax.scipy
function, the second addsjax.scipy.sqrtm
.