import torch

def layer_norm(x, gamma, beta, eps=1e-5):
    """x: [B, T, d]; gamma / beta: [d]."""
    z = x.to(torch.float32) if x.dtype in (
        torch.float16, torch.bfloat16) else x
    mu = z.mean(dim=-1, keepdim=True)
    centered = z - mu
    sigma_sq = centered.square().mean(dim=-1, keepdim=True)
    normalized = centered * torch.rsqrt(sigma_sq + eps)
    y = normalized * gamma + beta
    return y.to(x.dtype)

import torch

def rms_norm(x, gamma, eps=1e-5):
    """x: [B, T, d]; gamma: [d]."""
    z = x.to(torch.float32) if x.dtype in (
        torch.float16, torch.bfloat16) else x
    mean_square = z.square().mean(dim=-1, keepdim=True)
    normalized = z * torch.rsqrt(mean_square + eps)
    y = normalized * gamma
    return y.to(x.dtype)

import torch

def l2_normalize(x, eps=1e-5):
    """x: [..., d]; each final-axis vector independently."""
    z = x.to(torch.float32) if x.dtype in (
        torch.float16, torch.bfloat16) else x
    length = torch.linalg.vector_norm(z, dim=-1, keepdim=True)
    normalized = z / length.clamp_min(eps)
    return normalized.to(x.dtype)

import torch
from torch import nn

class AdaLN(nn.Module):
    def __init__(self, d, condition_dim):
        super().__init__()
        self.norm = nn.LayerNorm(d, eps=1e-5, elementwise_affine=False)
        self.modulation = nn.Sequential(
            nn.SiLU(), nn.Linear(condition_dim, 2 * d))
        nn.init.zeros_(self.modulation[-1].weight)
        nn.init.zeros_(self.modulation[-1].bias)

    def forward(self, x, condition):
        """x: [B, T, d]; condition: [B, condition_dim]."""
        shift, scale = self.modulation(condition).chunk(2, dim=-1)
        normalized = self.norm(x)
        return normalized * (1 + scale[:, None, :]) + shift[:, None, :]