# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/utils/03_normalizations.ipynb.

# %% auto 0
__all__ = ['LayerNorm2D', 'LayerNorm3D', 'DyT', 'get_norm_layer']

# %% ../../nbs/utils/03_normalizations.ipynb 2
import torch
from einops import rearrange
from torch import nn

# %% ../../nbs/utils/03_normalizations.ipynb 5
class LayerNorm2D(nn.LayerNorm):
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        input = rearrange(input, "b c y x -> b y x c").contiguous()
        input = super().forward(input)
        input = rearrange(input, "b y x c -> b c y x").contiguous()
        return input

# %% ../../nbs/utils/03_normalizations.ipynb 7
class LayerNorm3D(nn.LayerNorm):
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        input = rearrange(input, "b c z y x -> b z y x c").contiguous()
        input = super().forward(input)
        input = rearrange(input, "b z y x c -> b c z y x").contiguous()
        return input

# %% ../../nbs/utils/03_normalizations.ipynb 9
class DyT(nn.Module):
    # As proposed in Transformers without Normalization: https://arxiv.org/pdf/2503.10622

    def __init__(self, normalized_shape: int | list[int], alpha0: float = 0.5):
        super().__init__()

        if isinstance(normalized_shape, int):
            normalized_shape = (normalized_shape,)

        self.normalized_shape = normalized_shape
        self.alpha0 = alpha0

        self.alpha = nn.Parameter(torch.tensor([alpha0]))
        self.weight = nn.Parameter(torch.ones(normalized_shape, dtype=torch.float32))
        self.bias = nn.Parameter(torch.zeros(normalized_shape, dtype=torch.float32))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = torch.tanh(self.alpha * x)
        x = x * self.weight + self.bias
        return x

    def extra_repr(self) -> str:
        return f"normalized_shape={self.normalized_shape}, alpha0={self.alpha0}"

# %% ../../nbs/utils/03_normalizations.ipynb 12
def get_norm_layer(normalization_name: str, *args, **kwargs):
    if normalization_name is None:
        norm_layer = nn.Identity()
    elif normalization_name in {"layernorm", "layernorm1d"}:
        norm_layer = nn.LayerNorm(*args, **kwargs)
    elif normalization_name == "layernorm2d":
        norm_layer = LayerNorm2D(*args, **kwargs)
    elif normalization_name == "layernorm3d":
        norm_layer = LayerNorm3D(*args, **kwargs)
    elif normalization_name in {"batchnorm", "batchnorm1d"}:
        norm_layer = nn.BatchNorm1d(*args, **kwargs)
    elif normalization_name == "batchnorm2d":
        norm_layer = nn.BatchNorm2d(*args, **kwargs)
    elif normalization_name == "batchnorm3d":
        norm_layer = nn.BatchNorm3d(*args, **kwargs)
    elif normalization_name == "instancenorm1d":
        norm_layer = nn.InstanceNorm1d(*args, **kwargs)
    elif normalization_name == "instancenorm2d":
        norm_layer = nn.InstanceNorm2d(*args, **kwargs)
    elif normalization_name == "instancenorm3d":
        norm_layer = nn.InstanceNorm3d(*args, **kwargs)
    elif normalization_name == "groupnorm":
        norm_layer = nn.GroupNorm(*args, **kwargs)
    elif normalization_name == "dyt":
        norm_layer = DyT(*args, **kwargs)
    else:
        raise ValueError(f"Normalization {normalization_name} not implemented")

    return norm_layer
