def control_points(mean, cov):
  """Construct "sigma points" using a matrix sqrt (Cholesky or SVD are fine)."""
  sqrtm_cov = sqrtm(cov)  # or could be jax.scipy.linalg.cholesky(cov)
  offsets = jnp.sqrt(mean.shape[-1] + 0.5) * jnp.concatenate(
      [jnp.zeros_like(mean[..., None]), sqrtm_cov, -sqrtm_cov], -1)
  return mean[..., None] + offsets
Construct "sigma points" using a matrix sqrt (Cholesky or SVD are fine).
internal/mip_test.py