import torch


def covariance_z_mean(z_mean):
    """Computes the covariance of z_mean.
    Borrowed from https://github.com/google-research/disentanglement_lib/
    Uses cov(z_mean) = E[z_mean*z_mean^T] - E[z_mean]E[z_mean]^T.
    Args:
      z_mean: Encoder mean, tensor of size [batch_size, num_latent].
    Returns:
      cov_z_mean: Covariance of encoder mean, tensor of size [num_latent, num_latent].
    """
    expectation_z_mean_z_mean_t = torch.mean(z_mean.unsqueeze(2) * z_mean.unsqueeze(1), dim=0)
    expectation_z_mean = torch.mean(z_mean, dim=0)
    cov_z_mean = expectation_z_mean_z_mean_t - \
                 (expectation_z_mean.unsqueeze(1) * expectation_z_mean.unsqueeze(0))
    return cov_z_mean