An implementation of jax.scipy.linalg.sqrtm that is differentiable

[Newtech66]

Based on Denman-Beavers iterations as given in Lemma 3.2 of Cheng et al. (2001). I found the reference in Higham (2010).

def db_iter_sqrt(I, X, *args):
    X1i = jnp.linalg.inv(X[1])
    return ((0.5 * X[0] @ (I + X1i), 0.5 * (I + 0.5 * (X[1] + X1i))), None)
def sqrtm(A, s=10):
    return jax.lax.scan(partial(db_iter_sqrt, jnp.eye(A.shape[0])), (A, A), length=s)[0][0]
sqrtm_jit = jax.jit(sqrtm)

A possible workaround for #16579. Tagging @jakevdp in case they are interested (if not, my apologies for the unnecessary tag).

I used %timeit to measure the timings and checked the accuracy wrt jax.scipy.linalg.sqrtm on my machine using a CPU backend. The performance hit is not too great. And the accuracy is quite good too.

[jakevdp]

Hey - thanks for bringing this up! This is really interesting – the only thing that gives me pause here is the explicit inverse computation, which I suspect may cause the algorithm to be unstable or inaccurate for near-singular matrices. Have you looked at the accuracy in that domain?

[Newtech66]

I don't know how I would randomly generate near-singular matrices. But following this paper, which states that Hilbert matrices (as a special case of Hankel matrices) have a very high condition number, I tried the following:

A = osp.linalg.hilbert(n)
print(onp.linalg.cond(A))
Asqrt = sqrtm(A)
print(jnp.max(jnp.abs(Asqrt @ Asqrt - A)))
Asqrt = osp.linalg.sqrtm(A)
print(jnp.max(jnp.abs(Asqrt @ Asqrt - A)))

and indeed, the results grow worse. For n=100, the largest absolute entrywise error was ~60.0, with a condition number of ~3.3e+19. The currently implmented sqrt method remains stable with an error of ~2e-15. So this solution is not ideal for near-singular matrices.

[Newtech66]

Another option is to use Newton-Schulz iterations as described here. But it needs a norm computation and is slightly less accurate. But it works for the Hilbert matrices. (Correction: only for the square root, and not for the inverse square root it seems...)

def sqrt_isqrtm2(A):
    A_f = jnp.linalg.matrix_norm(A, ord='fro')
    A /= A_f
    s = 30
    I = jnp.eye(A.shape[0], dtype=A.dtype)
    def ns_iter_sqrt(X, _):
        ZY = 0.5 * (3 * I - X[1] @ X[0])
        return ((X[0] @ ZY, ZY @ X[1]), None)
    Asqrt, Aisqrt = jax.lax.scan(db_iter_sqrt, (A, I), length=s)[0][:2]
    return jnp.sqrt(A_f) * Asqrt, Aisqrt / jnp.sqrt(A_f)
sqrt_isqrtm2 = jax.jit(sqrt_isqrtm2)

I should note that both of these methods only work for positive definite matrices (all eigenvalues are positive). This works good enough for my use case, but not in the general case.