import torch


class LayerNorm2d(torch.nn.Module):
    def __init__(self, num_channels: int, epsilon: float = 1e-6):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.ones(num_channels))
        self.bias = torch.nn.Parameter(torch.zeros(num_channels))
        self.epsilon = epsilon

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.dim() != 4:
            raise Exception("LayerNorm2d input must be four dimensional.")

        mean = x.mean(dim=1, keepdim=True)
        variance = (x - mean).pow(2).mean(dim=1, keepdim=True)
        x = (x - mean) / torch.sqrt(variance + self.epsilon)
        x = x * self.weight.view(1, -1, 1, 1) + self.bias.view(1, -1, 1, 1)
        return x
