Add dispatcher for sqrtm

qutip / qutip-jax

JAX backend for QuTiP

BSD 3-Clause "New" or "Revised" License

15 stars 7 forks source link

Closed rochisha0 closed 3 weeks ago

rochisha0 commented 3 weeks ago

Description This PR aims to add dispatcher for sqrtm to jax.jit funtionality in qutip.core.metrics, a similar PR has also been in qutip

Note: jax.grad still does not work because there is no differentiation available in jax.grad for Schur decomposition

Ericgig commented 3 weeks ago

Test failure due to deprecation in the new jax version... New tests passes.