import torch


def compute_active_units(mu, delta):
    """Computes an estimate of the number of active units in the latent space.

    Args:
        mu(torch.FloatTensor): [n_samples, z_dim]. Batch of posterior means.
        delta(float): variance threshold. Latent dimensions with a variance above this threshold are active.

    Returns:
        int: the number of active dimensions.
    """
    outer_expectation = torch.mean(mu, 0) ** 2
    inner_expectation = torch.mean(mu ** 2, 0)
    return torch.sum(inner_expectation - outer_expectation > delta).item()