Function layer_norm_2d

tests/dump_phase7_reference.py:147–151

(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, eps: float = 1e-6)

def layer_norm_2d(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    u = x.mean(1, keepdim=True)
    s = (x - u).pow(2).mean(1, keepdim=True)
    x = (x - u) / torch.sqrt(s + eps)
    return weight[:, None, None] * x + bias[:, None, None]


def sinusoidal_pe(h: int, w: int, d_model: int) -> torch.Tensor: