torchembed

 20class RotaryEmbedding(nn.Module):
 21    """Rotary Position Embedding (RoPE).
 22
 23    Encodes position by rotating query and key vectors in 2D subspaces.
 24    Unlike additive embeddings, RoPE is applied directly to Q and K inside
 25    the attention layer, not to the input sequence.
 26
 27    Used in: LLaMA, Mistral, Falcon, PaLM, GPT-NeoX, and most modern LLMs.
 28
 29    Reference:
 30        Su et al., "RoFormer: Enhanced Transformer with Rotary Position Embedding"
 31        https://arxiv.org/abs/2104.09864
 32
 33    Args:
 34        dim: Head dimension (must be even). Typically d_model // num_heads.
 35        max_seq_len: Maximum sequence length to precompute. Longer sequences
 36            will be computed on the fly.
 37        base: Base for the geometric progression of frequencies. Default 10000
 38            matches the original paper. LLaMA 3 uses 500000.
 39        use_fused: If True, uses a fused triton kernel for the forward pass
 40            (requires GPU and ``triton``). Default False.
 41        device: Device to create buffers on.
 42
 43    Example::
 44
 45        rope = RotaryEmbedding(dim=64)
 46        q = torch.randn(2, 8, 16, 64)  # (batch, heads, seq, dim)
 47        k = torch.randn(2, 8, 16, 64)
 48        q_rot, k_rot = rope(q, k)
 49    """
 50
 51    def __init__(
 52        self,
 53        dim: int,
 54        max_seq_len: int = 2048,
 55        base: int = 10_000,
 56        use_fused: bool = False,
 57        device: Optional[torch.device] = None,
 58    ) -> None:
 59        super().__init__()
 60        if dim % 2 != 0:
 61            raise ValueError(f"dim must be even, got {dim}")
 62
 63        self.dim = dim
 64        self.max_seq_len = max_seq_len
 65        self.base = base
 66        self.use_fused = use_fused
 67
 68        # Precompute inverse frequencies: shape (dim/2,)
 69        inv_freq = 1.0 / (
 70            base ** (torch.arange(0, dim, 2, device=device).float() / dim)
 71        )  # noqa: E501
 72        self.register_buffer("inv_freq", inv_freq, persistent=False)
 73
 74        # Precompute cos/sin cache
 75        self._build_cache(max_seq_len, device)
 76
 77    def _build_cache(self, seq_len: int, device: Optional[torch.device] = None) -> None:
 78        t = torch.arange(seq_len, device=device or self.inv_freq.device).float()
 79        freqs = torch.outer(t, self.inv_freq)  # (seq_len, dim/2)
 80        emb = torch.cat([freqs, freqs], dim=-1)  # (seq_len, dim)
 81        self.register_buffer("cos_cache", emb.cos(), persistent=False)
 82        self.register_buffer("sin_cache", emb.sin(), persistent=False)
 83
 84    @staticmethod
 85    def _rotate_half(x: Tensor) -> Tensor:
 86        """Rotate the last dimension by splitting and negating halves."""
 87        x1, x2 = x.chunk(2, dim=-1)
 88        return torch.cat([-x2, x1], dim=-1)
 89
 90    def _vanilla_forward(
 91        self, q: Tensor, k: Tensor, cos: Tensor, sin: Tensor
 92    ) -> tuple[Tensor, Tensor]:  # noqa: E501
 93        cos = cos.unsqueeze(0).unsqueeze(0)
 94        sin = sin.unsqueeze(0).unsqueeze(0)
 95        q_rot = q * cos + self._rotate_half(q) * sin
 96        k_rot = k * cos + self._rotate_half(k) * sin
 97        return q_rot, k_rot
 98
 99    def _fused_forward(
100        self, q: Tensor, k: Tensor, cos: Tensor, sin: Tensor
101    ) -> tuple[Tensor, Tensor]:  # noqa: E501
102        from torchembed._triton import fused_rope_forward
103
104        return fused_rope_forward(q, k, cos, sin)
105
106    def forward(self, q: Tensor, k: Tensor, seq_dim: int = -2) -> tuple[Tensor, Tensor]:
107        """Apply rotary embeddings to query and key tensors.
108
109        Args:
110            q: Query tensor of shape (..., seq_len, dim).
111            k: Key tensor of shape (..., seq_len, dim).
112            seq_dim: Dimension along which sequence length lives. Default -2.
113
114        Returns:
115            Tuple of (rotated_q, rotated_k) with the same shapes as inputs.
116        """
117        seq_len = q.shape[seq_dim]
118
119        if seq_len > self.max_seq_len:
120            self._build_cache(seq_len, q.device)
121            self.max_seq_len = seq_len
122
123        cos = self.cos_cache[:seq_len].to(device=q.device)
124        sin = self.sin_cache[:seq_len].to(device=q.device)
125
126        if self.use_fused and q.is_cuda and k.is_cuda:
127            try:
128                return self._fused_forward(q, k, cos, sin)
129            except (ImportError, RuntimeError):
130                pass
131
132        return self._vanilla_forward(q, k, cos, sin)

135class ALiBiEmbedding(nn.Module):
136    """Attention with Linear Biases (ALiBi).
137
138    Instead of adding positional information to token embeddings, ALiBi
139    adds a fixed, non-learned bias to attention scores that penalizes
140    distance between tokens linearly. This allows strong extrapolation
141    to longer sequences than seen during training.
142
143    Used in: BLOOM, MPT, and other long-context models.
144
145    Reference:
146        Press et al., "Train Short, Test Long: Attention with Linear Biases
147        Enables Input Length Extrapolation" https://arxiv.org/abs/2108.12409
148
149    Args:
150        num_heads: Number of attention heads. Each head gets a different slope.
151        max_seq_len: Maximum sequence length to precompute biases for.
152
153    Example::
154
155        alibi = ALiBiEmbedding(num_heads=8)
156        # attn_scores: (batch, heads, seq, seq)
157        attn_scores = torch.randn(2, 8, 16, 16)
158        biased_scores = alibi(attn_scores)
159    """
160
161    def __init__(self, num_heads: int, max_seq_len: int = 2048) -> None:
162        super().__init__()
163        self.num_heads = num_heads
164
165        slopes = self._get_slopes(num_heads)  # (num_heads,)
166        bias = self._build_bias(slopes, max_seq_len)  # (num_heads, seq, seq)
167        self.register_buffer("bias", bias, persistent=False)
168
169    @staticmethod
170    def _get_slopes(num_heads: int) -> Tensor:
171        """Compute ALiBi slopes following the original paper's geometric sequence."""
172        # Nearest power of 2 >= num_heads
173        n = 2 ** math.ceil(math.log2(num_heads))
174        slopes = torch.pow(2, -torch.arange(1, n + 1) * (8 / n))
175        if n > num_heads:
176            # Interleave to handle non-power-of-2 head counts
177            slopes = torch.cat([slopes[1::2], slopes[::2]])[:num_heads]
178        return slopes
179
180    @staticmethod
181    def _build_bias(slopes: Tensor, max_seq_len: int) -> Tensor:
182        positions = torch.arange(max_seq_len)
183        # Relative distances: (seq, seq) lower-triangular distance matrix
184        dist = positions.unsqueeze(0) - positions.unsqueeze(1)  # (seq, seq)
185        dist = -dist.abs()
186        # Scale by each head's slope: (num_heads, seq, seq)
187        bias = slopes.unsqueeze(-1).unsqueeze(-1) * dist.unsqueeze(0)
188        return bias
189
190    def forward(self, attn_scores: Tensor) -> Tensor:
191        """Add ALiBi positional bias to attention scores.
192
193        Args:
194            attn_scores: Attention logits of shape (batch, heads, seq_q, seq_k).
195
196        Returns:
197            Attention scores with ALiBi bias added, same shape as input.
198        """
199        seq_len = attn_scores.shape[-1]
200        bias = self.bias[:, :seq_len, :seq_len]  # (heads, seq, seq)
201        return attn_scores + bias.unsqueeze(0)  # broadcast over batch

204class SinusoidalEmbedding(nn.Module):
205    """Fixed sinusoidal positional embedding from "Attention Is All You Need".
206
207    Adds a non-learned, frequency-based positional signal to input embeddings.
208    The encoding is deterministic and can generalize slightly beyond the
209    training sequence length.
210
211    Reference:
212        Vaswani et al., "Attention Is All You Need" https://arxiv.org/abs/1706.03762
213
214    Args:
215        dim: Embedding dimension (must be even).
216        max_seq_len: Maximum supported sequence length.
217        dropout: Optional dropout rate applied after adding the embedding.
218        learned_scale: If True, adds a single learned scalar to scale the
219            sinusoidal signal (a light touch of trainability).
220
221    Example::
222
223        emb = SinusoidalEmbedding(dim=512)
224        x = torch.randn(2, 16, 512)   # (batch, seq, dim)
225        x = emb(x)
226    """
227
228    def __init__(
229        self,
230        dim: int,
231        max_seq_len: int = 4096,
232        dropout: float = 0.0,
233        learned_scale: bool = False,
234    ) -> None:
235        super().__init__()
236        if dim % 2 != 0:
237            raise ValueError(f"dim must be even, got {dim}")
238
239        self.dim = dim
240        self.dropout = nn.Dropout(p=dropout) if dropout > 0 else nn.Identity()
241        self.scale = nn.Parameter(torch.ones(1)) if learned_scale else None
242
243        pe = self._build_pe(dim, max_seq_len)  # (1, max_seq_len, dim)
244        self.register_buffer("pe", pe, persistent=False)
245
246    @staticmethod
247    def _build_pe(dim: int, max_seq_len: int) -> Tensor:
248        position = torch.arange(max_seq_len).unsqueeze(1).float()
249        div_term = torch.exp(
250            torch.arange(0, dim, 2).float() * (-math.log(10000.0) / dim)
251        )
252        pe = torch.zeros(1, max_seq_len, dim)
253        pe[0, :, 0::2] = torch.sin(position * div_term)
254        pe[0, :, 1::2] = torch.cos(position * div_term)
255        return pe
256
257    def forward(self, x: Tensor) -> Tensor:
258        """Add sinusoidal positional encoding to input.
259
260        Args:
261            x: Input tensor of shape (batch, seq_len, dim).
262
263        Returns:
264            Tensor of same shape with positional encoding added.
265        """
266        seq_len = x.shape[1]
267        pe = self.pe[:, :seq_len, :]
268        if self.scale is not None:
269            pe = pe * self.scale
270        return self.dropout(x + pe)

273class LearnedPositionalEmbedding(nn.Module):
274    """Standard learned positional embedding.
275
276    A simple lookup table mapping each position index to a learnable vector.
277    Used in BERT, GPT-2, and many other models. Simpler than sinusoidal but
278    cannot extrapolate beyond the training sequence length.
279
280    Args:
281        max_seq_len: Maximum sequence length (vocabulary size for positions).
282        dim: Embedding dimension.
283        dropout: Optional dropout rate.
284        padding_idx: If set, the embedding at this index is not updated.
285
286    Example::
287
288        emb = LearnedPositionalEmbedding(max_seq_len=512, dim=768)
289        x = torch.randn(2, 16, 768)
290        x = emb(x)
291    """
292
293    def __init__(
294        self,
295        max_seq_len: int,
296        dim: int,
297        dropout: float = 0.0,
298        padding_idx: Optional[int] = None,
299    ) -> None:
300        super().__init__()
301        self.embedding = nn.Embedding(max_seq_len, dim, padding_idx=padding_idx)
302        self.dropout = nn.Dropout(p=dropout) if dropout > 0 else nn.Identity()
303        nn.init.normal_(self.embedding.weight, std=0.02)
304
305    def forward(self, x: Tensor, offset: int = 0) -> Tensor:
306        """Add learned positional embeddings to input.
307
308        Args:
309            x: Input tensor of shape (batch, seq_len, dim).
310            offset: Starting position index. Useful for KV-cache inference
311                where you process one token at a time.
312
313        Returns:
314            Tensor of same shape with positional embeddings added.
315        """
316        seq_len = x.shape[1]
317        positions = torch.arange(offset, offset + seq_len, device=x.device)
318        return self.dropout(x + self.embedding(positions))

18class RandomFourierFeatures(nn.Module):
19    """Random Fourier Features (RFF) for kernel approximation.
20
21    Maps input features to a randomized low-dimensional feature space such
22    that the inner product in the new space approximates a shift-invariant
23    kernel (e.g. RBF/Gaussian). Useful for scalable kernel methods and as
24    a fixed encoding for continuous inputs like coordinates.
25
26    Reference:
27        Rahimi & Recht, "Random Features for Large-Scale Kernel Machines"
28        https://papers.nips.cc/paper/2007/hash/013a006f03dbc5392effeb8f18fda755-Abstract.html
29
30    Args:
31        in_features: Input dimension.
32        out_features: Output dimension (number of random features). Should be
33            large enough to approximate the kernel well (e.g. 256–2048).
34        sigma: Bandwidth of the RBF kernel. Controls the frequency of the
35            random features. Smaller = higher frequency.
36        trainable: If True, the random projection matrix is a learnable parameter.
37
38    Example::
39
40        rff = RandomFourierFeatures(in_features=2, out_features=256)
41        coords = torch.randn(32, 2)   # e.g. 2D spatial coordinates
42        features = rff(coords)        # (32, 256)
43    """
44
45    def __init__(
46        self,
47        in_features: int,
48        out_features: int,
49        sigma: float = 1.0,
50        trainable: bool = False,
51    ) -> None:
52        super().__init__()
53        if out_features % 2 != 0:
54            raise ValueError(f"out_features must be even, got {out_features}")
55
56        self.in_features = in_features
57        self.out_features = out_features
58        self.sigma = sigma
59
60        # Random projection matrix: (in_features, out_features // 2)
61        W = torch.randn(in_features, out_features // 2) / sigma
62        if trainable:
63            self.W = nn.Parameter(W)
64        else:
65            self.register_buffer("W", W)
66
67    def forward(self, x: Tensor) -> Tensor:
68        """Project input to random Fourier feature space.
69
70        Args:
71            x: Input tensor of shape (..., in_features).
72
73        Returns:
74            Tensor of shape (..., out_features) with cosine and sine features
75            concatenated and scaled by sqrt(2 / out_features).
76        """
77        projection = x @ self.W      # (..., out_features // 2)
78        scale = math.sqrt(2.0 / self.out_features)
79        return scale * torch.cat([torch.cos(projection), torch.sin(projection)], dim=-1)

 82class LearnedFourierFeatures(nn.Module):
 83    """Learned Fourier Features.
 84
 85    A trainable variant of random Fourier features where both the frequency
 86    matrix and a final linear projection are learned end-to-end. Useful when
 87    you want the network to discover the right frequency decomposition for
 88    your specific input domain.
 89
 90    Args:
 91        in_features: Input dimension.
 92        num_frequencies: Number of frequency components (output will be
 93            2 * num_frequencies before projection).
 94        out_features: Final output dimension after linear projection.
 95        sigma_init: Initial scale for frequency initialization.
 96
 97    Example::
 98
 99        lff = LearnedFourierFeatures(
100            in_features=3, num_frequencies=128, out_features=256
101        )
102        x = torch.randn(16, 3)
103        features = lff(x)    # (16, 256)
104    """
105
106    def __init__(
107        self,
108        in_features: int,
109        num_frequencies: int,
110        out_features: int,
111        sigma_init: float = 1.0,
112    ) -> None:
113        super().__init__()
114        self.freq = nn.Parameter(
115            torch.randn(in_features, num_frequencies) / sigma_init
116        )
117        self.proj = nn.Linear(2 * num_frequencies, out_features)
118        nn.init.xavier_uniform_(self.proj.weight)
119
120    def forward(self, x: Tensor) -> Tensor:
121        """Compute learned Fourier features.
122
123        Args:
124            x: Input tensor of shape (..., in_features).
125
126        Returns:
127            Tensor of shape (..., out_features).
128        """
129        projection = x @ self.freq    # (..., num_frequencies)
130        features = torch.cat([torch.cos(projection), torch.sin(projection)], dim=-1)
131        return self.proj(features)

134class GaussianFourierProjection(nn.Module):
135    """Gaussian Fourier Projection for continuous scalar embedding.
136
137    A widely used technique in diffusion models and score-based generative
138    models to embed continuous scalars (e.g. timestep t, noise level σ)
139    into a high-dimensional space. Uses fixed random frequencies scaled by
140    a learnable or fixed bandwidth.
141
142    Reference:
143        Song et al., "Score-Based Generative Modeling through Stochastic
144        Differential Equations" https://arxiv.org/abs/2011.13456
145
146    Args:
147        embed_dim: Output embedding dimension. Must be even.
148        scale: Scale factor for the random frequencies. Controls how quickly
149            the embedding oscillates. Typical values: 16–30.
150        learnable: If True, the random weights are trainable.
151
152    Example::
153
154        gfp = GaussianFourierProjection(embed_dim=256, scale=16)
155        t = torch.rand(32)        # continuous timesteps in [0, 1]
156        emb = gfp(t)              # (32, 256)
157
158        # Common usage in diffusion models:
159        # t_emb = gfp(t)
160        # then feed through an MLP to condition the UNet
161    """
162
163    def __init__(
164        self,
165        embed_dim: int,
166        scale: float = 16.0,
167        learnable: bool = False,
168    ) -> None:
169        super().__init__()
170        if embed_dim % 2 != 0:
171            raise ValueError(f"embed_dim must be even, got {embed_dim}")
172
173        W = torch.randn(embed_dim // 2) * scale
174        if learnable:
175            self.W = nn.Parameter(W)
176        else:
177            self.register_buffer("W", W)
178
179    def forward(self, x: Tensor) -> Tensor:
180        """Embed a continuous scalar input.
181
182        Args:
183            x: Scalar input tensor of shape (batch,) or (batch, 1).
184                Values are typically in [0, 1] or [0, T].
185
186        Returns:
187            Tensor of shape (batch, embed_dim).
188        """
189        if x.dim() == 1:
190            x = x.unsqueeze(-1)    # (batch, 1)
191        projection = x * self.W.unsqueeze(0) * 2 * math.pi   # (batch, embed_dim//2)
192        return torch.cat([torch.sin(projection), torch.cos(projection)], dim=-1)

29class EntityEmbedding(nn.Module):
30    """Learned embedding for a single categorical feature.
31
32    Wraps ``nn.Embedding`` with automatic dimension sizing and optional
33    dropout. The auto-sizing heuristic (fast.ai rule) empirically works well
34    for tabular data without manual tuning.
35
36    Reference:
37        Howard & Gugger, "Deep Learning for Coders with fastai and PyTorch"
38        https://arxiv.org/abs/2002.04688 (entity embeddings section)
39
40    Args:
41        num_categories: Vocabulary size (number of unique categories + 1 for
42            unknown). Indices must be in [0, num_categories - 1].
43        embed_dim: Embedding dimension. If None, uses the auto-sizing
44            heuristic: ``min(600, round(1.6 * num_categories ** 0.56))``.
45        dropout: Dropout rate applied to embeddings.
46        padding_idx: Index that will always produce a zero embedding
47            (e.g. for padding / unknown tokens).
48
49    Example::
50
51        emb = EntityEmbedding(num_categories=50)
52        x = torch.randint(0, 50, (32,))    # batch of category indices
53        out = emb(x)                        # (32, embed_dim)
54    """
55
56    def __init__(
57        self,
58        num_categories: int,
59        embed_dim: Optional[int] = None,
60        dropout: float = 0.0,
61        padding_idx: Optional[int] = None,
62    ) -> None:
63        super().__init__()
64        self.num_categories = num_categories
65        self.embed_dim = embed_dim or _auto_dim(num_categories)
66
67        self.embedding = nn.Embedding(
68            num_categories, self.embed_dim, padding_idx=padding_idx
69        )
70        self.dropout = nn.Dropout(p=dropout) if dropout > 0 else nn.Identity()
71
72        nn.init.normal_(self.embedding.weight, std=0.01)
73        if padding_idx is not None:
74            nn.init.zeros_(self.embedding.weight[padding_idx])
75
76    def forward(self, x: Tensor) -> Tensor:
77        """Embed categorical indices.
78
79        Args:
80            x: Long tensor of category indices, shape (...,).
81
82        Returns:
83            Float tensor of shape (..., embed_dim).
84        """
85        return self.dropout(self.embedding(x))

 88class MultiCategoricalEmbedding(nn.Module):
 89    """Joint embedding for multiple categorical columns.
 90
 91    Intended for tabular data where each row has several categorical features
 92    (e.g. country, product category, day of week). Creates one
 93    ``EntityEmbedding`` per column and concatenates the results.
 94
 95    Args:
 96        cardinalities: Sequence of vocabulary sizes, one per categorical column.
 97            E.g. ``[50, 7, 12]`` for columns with 50, 7, and 12 categories.
 98        embed_dims: Per-column embedding dimensions. If None, uses the
 99            auto-sizing heuristic for each column.
100        dropout: Dropout rate applied to each embedding independently.
101        padding_idx: Shared padding index applied to all columns.
102
103    Properties:
104        output_dim: Total output dimension (sum of all embed_dims).
105
106    Example::
107
108        # 3 categorical columns: country (50 cats), weekday (7), month (12)
109        emb = MultiCategoricalEmbedding(cardinalities=[50, 7, 12])
110        x = torch.stack([
111            torch.randint(0, 50, (32,)),
112            torch.randint(0, 7, (32,)),
113            torch.randint(0, 12, (32,)),
114        ], dim=1)                             # (32, 3)
115        out = emb(x)                          # (32, output_dim)
116        print(emb.output_dim)
117    """
118
119    def __init__(
120        self,
121        cardinalities: Sequence[int],
122        embed_dims: Optional[Sequence[int]] = None,
123        dropout: float = 0.0,
124        padding_idx: Optional[int] = None,
125    ) -> None:
126        super().__init__()
127
128        if embed_dims is not None and len(embed_dims) != len(cardinalities):
129            raise ValueError(
130                f"embed_dims length ({len(embed_dims)}) must match "
131                f"cardinalities length ({len(cardinalities)})"
132            )
133
134        dims = embed_dims or [None] * len(cardinalities)   # type: ignore[list-item]
135        self.embeddings = nn.ModuleList([
136            EntityEmbedding(n, d, dropout, padding_idx)
137            for n, d in zip(cardinalities, dims)
138        ])
139        self.output_dim: int = sum(e.embed_dim for e in self.embeddings)  # type: ignore[union-attr]
140
141    def forward(self, x: Tensor) -> Tensor:
142        """Embed all categorical columns and concatenate.
143
144        Args:
145            x: Long tensor of shape (batch, num_columns) containing
146               category indices. Column order must match ``cardinalities``.
147
148        Returns:
149            Float tensor of shape (batch, output_dim).
150        """
151        parts = [emb(x[:, i]) for i, emb in enumerate(self.embeddings)]
152        return torch.cat(parts, dim=-1)
153
154    def embedding_dims(self) -> list[tuple[int, int]]:
155        """Return list of (num_categories, embed_dim) for each column."""
156        return [(e.num_categories, e.embed_dim) for e in self.embeddings]  # type: ignore[union-attr]