Zhengyuan Su (苏政渊)andClaude Opus 4.8
2d0b1d8a29
Signed-off-by: Zhengyuan Su <su.zhengyuan@u.nus.edu> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
29 lines
917 B
Python
29 lines
917 B
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import math
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
def swish(x: torch.Tensor) -> torch.Tensor:
|
|
return x * torch.sigmoid(x)
|
|
|
|
|
|
class SinusoidalPositionalEncoding(nn.Module):
|
|
def __init__(self, embedding_dim: int) -> None:
|
|
super().__init__()
|
|
self.embedding_dim = embedding_dim
|
|
|
|
def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
|
|
timesteps = timesteps.float()
|
|
B, T = timesteps.shape
|
|
device = timesteps.device
|
|
|
|
half_dim = self.embedding_dim // 2
|
|
exponent = -torch.arange(half_dim, dtype=torch.float, device=device) * (math.log(10000.0) / half_dim)
|
|
freqs = timesteps.unsqueeze(-1) * exponent.exp() # (B, T, half_dim)
|
|
enc = torch.cat([torch.sin(freqs), torch.cos(freqs)], dim=-1) # (B, T, embedding_dim)
|
|
return enc
|