import torch


def ATE(y: torch.Tensor, hat_y: torch.Tensor) -> torch.Tensor:
    """
    Average Treatment Effect.
    ATE measures what is the expected causal effect of the treatment across all individuals in the population.
    Args:
        y: expected outcome.
        hat_y: estimated outcome.
    """
    return torch.abs(
        torch.mean(y[:, 1] - y[:, 0]) - torch.mean(hat_y[:, 1] - hat_y[:, 0])
    )