def sqrtm(mat):
  """Take the matrix square root of a PSD matrix [..., d, d]."""
  eigval, eigvec = jax.scipy.linalg.eigh(mat)
  scaling = jnp.sqrt(jnp.maximum(0., eigval))[..., None, :]
  return math.matmul(eigvec * scaling, jnp.moveaxis(eigvec, -2, -1))
Take the matrix square root of a PSD matrix [..., d, d].
internal/mip_test.py