import torch


def concat_mean_stats(inputs):
    """Add mean statistics to discriminator input.

    GANs often run into mode collapse since the discriminator sees every
    sample in isolation. I.e., it cannot detect whether all samples in a batch
    do look alike.

    A simple way to allow the discriminator to have access to batch statistics
    is to simply concatenate the mean (across batch dimension) of all
    discriminator samples to each sample.

    Args:
        inputs: The input batch to the discriminator.

    Returns:
        The modified input batch.
    """
    stats = torch.mean(inputs, 0, keepdim=True)
    stats = stats.expand(inputs.size())
    return torch.cat([stats, inputs], dim=1)