Adds matrix sqrt

[SaturdayGenfo]

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:

for j in range(n):
    for i in range(j-1, -1, -1):
        s = 0
        for k in range(i+1, j):
            s += U[i, k] * U[k,i]
        num = T[i, j] - s
        denom = diag[i] + diag[j]
        if num == 0.0:
            U = U.at[i,j].set(0.0)
        else:
            U = U.at[i,j].set(num/denom)  

(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 adds jax.scipy.sqrtm.

[shoyer]

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.

[SaturdayGenfo]

I ran a few tests (in double precision) and it performs as well as the non-blocked version of scipy.linalg.sqrtm. But it is slower than the sylvester augmented version with default blocksize=64.

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

[hawkinsp]

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.

[SaturdayGenfo]

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 5000x5000. To avoid these loops, one option is to write a custom call to Eigen like they do in tensorflow. Only issue is that custom calls are not the easiest to write.

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 sqrtm worthwhile. What do you think ?

[shoyer]

This should raise an error if an invalid output is provided

[SaturdayGenfo]

I added a line to check the output argument schur.

[shoyer]

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.