An einops-style generalized normalization layer.
Project description
einorm
An einops-style generalized normalization layer.
Installation
You need torch
>= 1.13 or functorch
to be installed:
pip install einorm
Usage
from einorm import Einorm
# Equivalent to nn.LayerNorm(1024)
Einorm("b n d", "d", d=1024)
# Specify the dimensions and sizes to normalize along.
Einorm("a b c d e", "b d", b=3, d=4)
According to ViT-22B, normalizing query and key in a head-wise fashion can help stabilize the training dynamics. This can be achieved by providing additional grouping arguments to Einorm
:
Einorm("b h n d", "d", "h", h=16, d=64) # num_heads=16, head_dim=64
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
einorm-0.1.1.tar.gz
(3.8 kB
view hashes)