matrix square root function only implemented on cpu?

[xinglong-li]

I got the error message :

NotImplementedError: Schur decomposition is only implemented on the CPU backend.

when I try use the function: jax.scipy.linalg.sqatm in computing the matrix square root on my GPU

[jakevdp]

That's true - the Schur decomposition was added in #8033, but it only added a CPU implementation. Adding GPU or TPU support would involve writing a translation rule to target the appropriate GPU or TPU implementation.

[jakevdp]

I took a look through the cusolver docs (cusolver is what JAX lowers to for GPU translation rules for some other decompositions); it doesn't look like there's a schur implementation available, unfortunately.

[xinglong-li]

Got it.
Thanks for the clarification ;)