def surface_stats(points):
  """Get the sample mean and covariance matrix of a set of matrices [..., d]."""
  means = jnp.mean(points, -1)
  centered = points - means[..., None]
  covs = jnp.mean(centered[..., None, :, :] * centered[..., :, None, :], -1)
  return means, covs
Get the sample mean and covariance matrix of a set of matrices [..., d].
internal/mip_test.py