475 lines
21 KiB
Python
475 lines
21 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||
#
|
||
# This source code is licensed under the Apache License, Version 2.0
|
||
# found in the LICENSE file in the root directory of this source tree.
|
||
|
||
|
||
# Implementation of 2D Rotary Position Embeddings (RoPE).
|
||
|
||
# This module provides a clean implementation of 2D Rotary Position Embeddings,
|
||
# which extends the original RoPE concept to handle 2D spatial positions.
|
||
|
||
# Inspired by:
|
||
# https://github.com/meta-llama/codellama/blob/main/llama/model.py
|
||
# https://github.com/naver-ai/rope-vit
|
||
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from typing import Dict, Tuple
|
||
|
||
from typing import List, Optional, Tuple, Union
|
||
|
||
|
||
class PositionGetter:
|
||
"""Generates and caches 2D spatial positions for patches in a grid.
|
||
|
||
This class efficiently manages the generation of spatial coordinates for patches
|
||
in a 2D grid, caching results to avoid redundant computations.
|
||
|
||
Attributes:
|
||
position_cache: Dictionary storing precomputed position tensors for different
|
||
grid dimensions.
|
||
"""
|
||
|
||
def __init__(self):
|
||
"""Initializes the position generator with an empty cache."""
|
||
self.position_cache: Dict[Tuple[int, int], torch.Tensor] = {}
|
||
|
||
def __call__(self, batch_size: int, height: int, width: int, device: torch.device) -> torch.Tensor:
|
||
"""Generates spatial positions for a batch of patches.
|
||
|
||
Args:
|
||
batch_size: Number of samples in the batch.
|
||
height: Height of the grid in patches.
|
||
width: Width of the grid in patches.
|
||
device: Target device for the position tensor.
|
||
|
||
Returns:
|
||
Tensor of shape (batch_size, height*width, 2) containing y,x coordinates
|
||
for each position in the grid, repeated for each batch item.
|
||
"""
|
||
if (height, width) not in self.position_cache:
|
||
y_coords = torch.arange(height, device=device)
|
||
x_coords = torch.arange(width, device=device)
|
||
positions = torch.cartesian_prod(y_coords, x_coords)
|
||
self.position_cache[height, width] = positions
|
||
|
||
cached_positions = self.position_cache[height, width]
|
||
return cached_positions.view(1, height * width, 2).expand(batch_size, -1, -1).clone()
|
||
|
||
|
||
class RotaryPositionEmbedding2D(nn.Module):
|
||
"""2D Rotary Position Embedding implementation.
|
||
|
||
This module applies rotary position embeddings to input tokens based on their
|
||
2D spatial positions. It handles the position-dependent rotation of features
|
||
separately for vertical and horizontal dimensions.
|
||
|
||
Args:
|
||
frequency: Base frequency for the position embeddings. Default: 100.0
|
||
scaling_factor: Scaling factor for frequency computation. Default: 1.0
|
||
|
||
Attributes:
|
||
base_frequency: Base frequency for computing position embeddings.
|
||
scaling_factor: Factor to scale the computed frequencies.
|
||
frequency_cache: Cache for storing precomputed frequency components.
|
||
"""
|
||
|
||
def __init__(self, frequency: float = 100.0, scaling_factor: float = 1.0):
|
||
"""Initializes the 2D RoPE module."""
|
||
super().__init__()
|
||
self.base_frequency = frequency
|
||
self.scaling_factor = scaling_factor
|
||
self.frequency_cache: Dict[Tuple, Tuple[torch.Tensor, torch.Tensor]] = {}
|
||
|
||
def _compute_frequency_components(
|
||
self, dim: int, seq_len: int, device: torch.device, dtype: torch.dtype
|
||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||
"""Computes frequency components for rotary embeddings.
|
||
|
||
Args:
|
||
dim: Feature dimension (must be even).
|
||
seq_len: Maximum sequence length.
|
||
device: Target device for computations.
|
||
dtype: Data type for the computed tensors.
|
||
|
||
Returns:
|
||
Tuple of (cosine, sine) tensors for frequency components.
|
||
"""
|
||
cache_key = (dim, seq_len, device, dtype)
|
||
if cache_key not in self.frequency_cache:
|
||
# Compute frequency bands
|
||
exponents = torch.arange(0, dim, 2, device=device).float() / dim
|
||
inv_freq = 1.0 / (self.base_frequency**exponents)
|
||
|
||
# Generate position-dependent frequencies
|
||
positions = torch.arange(seq_len, device=device, dtype=inv_freq.dtype)
|
||
angles = torch.einsum("i,j->ij", positions, inv_freq)
|
||
|
||
# Compute and cache frequency components
|
||
angles = angles.to(dtype)
|
||
angles = torch.cat((angles, angles), dim=-1)
|
||
cos_components = angles.cos().to(dtype)
|
||
sin_components = angles.sin().to(dtype)
|
||
self.frequency_cache[cache_key] = (cos_components, sin_components)
|
||
|
||
return self.frequency_cache[cache_key]
|
||
|
||
@staticmethod
|
||
def _rotate_features(x: torch.Tensor) -> torch.Tensor:
|
||
"""Performs feature rotation by splitting and recombining feature dimensions.
|
||
|
||
Args:
|
||
x: Input tensor to rotate.
|
||
|
||
Returns:
|
||
Rotated feature tensor.
|
||
"""
|
||
feature_dim = x.shape[-1]
|
||
x1, x2 = x[..., : feature_dim // 2], x[..., feature_dim // 2 :]
|
||
return torch.cat((-x2, x1), dim=-1)
|
||
|
||
def _apply_1d_rope(
|
||
self, tokens: torch.Tensor, positions: torch.Tensor, cos_comp: torch.Tensor, sin_comp: torch.Tensor
|
||
) -> torch.Tensor:
|
||
"""Applies 1D rotary position embeddings along one dimension.
|
||
|
||
Args:
|
||
tokens: Input token features.
|
||
positions: Position indices.
|
||
cos_comp: Cosine components for rotation.
|
||
sin_comp: Sine components for rotation.
|
||
|
||
Returns:
|
||
Tokens with applied rotary position embeddings.
|
||
"""
|
||
# Embed positions with frequency components
|
||
cos = F.embedding(positions, cos_comp)[:, None, :, :]
|
||
sin = F.embedding(positions, sin_comp)[:, None, :, :]
|
||
|
||
# Apply rotation
|
||
return (tokens * cos) + (self._rotate_features(tokens) * sin)
|
||
|
||
def forward(self, tokens: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
|
||
"""Applies 2D rotary position embeddings to input tokens.
|
||
|
||
Args:
|
||
tokens: Input tensor of shape (batch_size, n_heads, n_tokens, dim).
|
||
The feature dimension (dim) must be divisible by 4.
|
||
positions: Position tensor of shape (batch_size, n_tokens, 2) containing
|
||
the y and x coordinates for each token.
|
||
|
||
Returns:
|
||
Tensor of same shape as input with applied 2D rotary position embeddings.
|
||
|
||
Raises:
|
||
AssertionError: If input dimensions are invalid or positions are malformed.
|
||
"""
|
||
# Validate inputs
|
||
assert tokens.size(-1) % 2 == 0, "Feature dimension must be even"
|
||
assert positions.ndim == 3 and positions.shape[-1] == 2, "Positions must have shape (batch_size, n_tokens, 2)"
|
||
|
||
# Compute feature dimension for each spatial direction
|
||
feature_dim = tokens.size(-1) // 2
|
||
|
||
# Get frequency components
|
||
max_position = int(positions.max()) + 1
|
||
cos_comp, sin_comp = self._compute_frequency_components(feature_dim, max_position, tokens.device, tokens.dtype)
|
||
|
||
# Split features for vertical and horizontal processing
|
||
vertical_features, horizontal_features = tokens.chunk(2, dim=-1)
|
||
|
||
# Apply RoPE separately for each dimension
|
||
vertical_features = self._apply_1d_rope(vertical_features, positions[..., 0], cos_comp, sin_comp)
|
||
horizontal_features = self._apply_1d_rope(horizontal_features, positions[..., 1], cos_comp, sin_comp)
|
||
|
||
# Combine processed features
|
||
return torch.cat((vertical_features, horizontal_features), dim=-1)
|
||
|
||
|
||
|
||
def get_1d_rotary_pos_embed(
|
||
dim: int,
|
||
pos: Union[np.ndarray, int],
|
||
theta: float = 10000.0,
|
||
use_real=False,
|
||
linear_factor=1.0,
|
||
ntk_factor=1.0,
|
||
repeat_interleave_real=True,
|
||
freqs_dtype=torch.float32, # torch.float32, torch.float64 (flux)
|
||
):
|
||
"""
|
||
计算1D旋转位置编码（RoPE）的频率张量。
|
||
|
||
RoPE的核心思想：使用旋转矩阵来编码位置信息，使得相对位置关系保持不变。
|
||
公式：对于位置m和维度i，频率为 θ_i = θ^(-2i/d)，其中θ是基础频率（默认10000）
|
||
|
||
Args:
|
||
dim: 特征维度，必须是偶数（因为要成对处理）
|
||
pos: 位置索引，可以是整数（自动生成0到pos-1的序列）或位置数组 [S]
|
||
theta: 基础频率，控制位置编码的周期性（默认10000）
|
||
use_real: 是否返回实数形式（cos和sin分开）还是复数形式
|
||
linear_factor: 线性缩放因子，用于上下文扩展
|
||
ntk_factor: NTK-Aware缩放因子，用于处理更长的序列
|
||
repeat_interleave_real: 当use_real=True时，是否交错重复（用于某些模型架构）
|
||
freqs_dtype: 频率张量的数据类型
|
||
|
||
Returns:
|
||
复数形式：[S, D/2] 的复数张量，表示 e^(i*m*θ_j)
|
||
实数形式：两个 [S, D] 的张量（cos和sin）
|
||
"""
|
||
# 确保维度是偶数（RoPE需要成对处理维度）
|
||
assert dim % 2 == 0
|
||
|
||
# 将位置转换为torch张量
|
||
if isinstance(pos, int):
|
||
pos = torch.arange(pos) # 生成 [0, 1, 2, ..., pos-1]
|
||
if isinstance(pos, np.ndarray):
|
||
pos = torch.from_numpy(pos) # [S]
|
||
|
||
# 应用NTK缩放（Neural Tangent Kernel，用于处理训练时未见过的长序列）
|
||
theta = theta * ntk_factor
|
||
|
||
# 步骤1：计算频率 θ_i = 1 / (θ^(2i/d))
|
||
# 其中 i ∈ {0, 2, 4, ..., dim-2}（只取偶数索引，因为成对处理）
|
||
# 公式：freq_i = 1 / (theta^(2i/d) * linear_factor)
|
||
freqs = (
|
||
1.0
|
||
/ (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype, device=pos.device)[: (dim // 2)] / dim))
|
||
/ linear_factor
|
||
) # [D/2]，每个频率对应一个维度对
|
||
|
||
# 步骤2：计算位置-频率矩阵
|
||
# 使用外积：pos[m] * freqs[i] = m * θ_i
|
||
# 结果：每个位置m和每个频率i的组合
|
||
freqs = torch.outer(pos, freqs) # [S, D/2]
|
||
|
||
# 步骤3：根据返回格式转换
|
||
if use_real and repeat_interleave_real:
|
||
# 方式1：交错重复（用于flux, hunyuan-dit, cogvideox等模型）
|
||
# 将每个频率的cos和sin交错排列：[cos_0, cos_0, cos_1, cos_1, ...]
|
||
freqs_cos = freqs.cos().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float() # [S, D]
|
||
freqs_sin = freqs.sin().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float() # [S, D]
|
||
return freqs_cos, freqs_sin
|
||
elif use_real:
|
||
# 方式2：拼接重复（用于stable audio, allegro等模型）
|
||
# 将所有cos拼接，然后是所有sin：[cos_0, cos_1, ..., cos_n, cos_0, cos_1, ..., cos_n]
|
||
freqs_cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).float() # [S, D]
|
||
freqs_sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).float() # [S, D]
|
||
return freqs_cos, freqs_sin
|
||
else:
|
||
# 方式3：复数形式（用于lumina等模型）
|
||
# 使用欧拉公式：e^(iθ) = cos(θ) + i*sin(θ)
|
||
# torch.polar(r, θ) 返回 r * e^(iθ)，这里r=1，所以就是 e^(i*freqs)
|
||
freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64: [S, D/2]
|
||
return freqs_cis
|
||
|
||
|
||
class WanRotaryPosEmbed(nn.Module):
|
||
"""
|
||
3D旋转位置编码（3D RoPE）模块
|
||
|
||
核心思想：将RoPE扩展到3D空间（时间、高度、宽度），为视频或3D数据提供位置编码。
|
||
每个维度（t, h, w）独立使用RoPE，然后拼接起来。
|
||
|
||
公式：
|
||
对于3D位置 (f, h, w)（帧、高度、宽度）：
|
||
- 帧维度使用 dim_f 个特征维度
|
||
- 高度维度使用 dim_h 个特征维度
|
||
- 宽度维度使用 dim_w 个特征维度
|
||
其中 dim_f + dim_h + dim_w = attention_head_dim
|
||
"""
|
||
def __init__(
|
||
self,
|
||
attention_head_dim: int,
|
||
patch_size: Tuple[int, int, int],
|
||
max_seq_len: int = 1024,
|
||
theta: float = 10000.0,
|
||
fhw_dim: Optional[Tuple[int, int, int]] = [20, 22, 22],
|
||
):
|
||
super().__init__()
|
||
|
||
self.attention_head_dim = attention_head_dim # 注意力头的总维度
|
||
self.patch_size = patch_size # patch大小 (patch_f, patch_h, patch_w)
|
||
self.max_seq_len = max_seq_len # 最大序列长度（用于预计算频率）
|
||
|
||
# 步骤1：分配维度给三个空间维度
|
||
if fhw_dim is not None:
|
||
# 如果指定了维度分配，使用指定的
|
||
assert attention_head_dim == sum(
|
||
fhw_dim
|
||
), f"attention_head_dim {attention_head_dim} must match sum(fhw_dim) {sum(fhw_dim)}"
|
||
t_dim, h_dim, w_dim = fhw_dim
|
||
else:
|
||
# 否则自动分配：h和w各占1/3，t占剩余
|
||
# 例如：如果attention_head_dim=64，则 h_dim=w_dim=21，t_dim=22
|
||
h_dim = w_dim = 2 * (attention_head_dim // 6)
|
||
t_dim = attention_head_dim - h_dim - w_dim
|
||
|
||
# 保存维度分配以便在forward中使用
|
||
self.fhw_dim = (t_dim, h_dim, w_dim)
|
||
|
||
# 步骤2：为每个维度预计算频率
|
||
# 分别计算时间、高度、宽度三个维度的RoPE频率
|
||
freqs = []
|
||
for dim in [t_dim, h_dim, w_dim]:
|
||
# 每个维度独立调用1D RoPE
|
||
# 返回复数形式的频率: [max_seq_len, dim//2]
|
||
freq = get_1d_rotary_pos_embed(
|
||
dim, max_seq_len, theta, use_real=False, repeat_interleave_real=False, freqs_dtype=torch.float64
|
||
)
|
||
freqs.append(freq)
|
||
# 将三个维度的频率在最后一维拼接: [max_seq_len, (t_dim + h_dim + w_dim)//2]
|
||
self.freqs = torch.cat(freqs, dim=1)
|
||
|
||
def forward(self, ppf, pph, ppw, patch_start_idx, device: torch.device, f_start: int = 0, f_end: Optional[int] = None) -> torch.Tensor:
|
||
"""
|
||
前向传播：为3D输入（视频帧+patch）生成旋转位置编码
|
||
|
||
参数：
|
||
- ppf (int): 帧数（patches per frame），当f_end为None时使用
|
||
- pph (int): 每帧的patch高度数量
|
||
- ppw (int): 每帧的patch宽度数量
|
||
- patch_start_idx (int): 每帧的特殊token数量（在patches之前）
|
||
- device: 计算设备（CPU/GPU）
|
||
- f_start (int): 起始帧索引（用于causal模式），默认为0
|
||
- f_end (Optional[int]): 结束帧索引（用于causal模式），如果为None则使用ppf作为帧数
|
||
|
||
返回：
|
||
- freqs: [1, 1, ppf * (patch_start_idx + pph * ppw), head_dim//2] 复数频率tensor
|
||
|
||
Token排列顺序：
|
||
[frame0_special_token_0, ..., frame0_special_token_N,
|
||
frame0_patch_0, ..., frame0_patch_M,
|
||
frame1_special_token_0, ..., frame1_special_token_N,
|
||
frame1_patch_0, ..., frame1_patch_M,
|
||
...]
|
||
|
||
模式：
|
||
- 非causal模式：f_end=None，使用ppf作为帧数，从位置0开始
|
||
- Causal模式：f_end不为None，使用[f_start, f_end)范围的帧，ppf会被重新计算
|
||
"""
|
||
|
||
# 步骤1：将预计算的频率移到目标设备，并分割成三个维度
|
||
self.freqs = self.freqs.to(device)
|
||
# 获取实际的维度分配
|
||
if hasattr(self, 'fhw_dim') and self.fhw_dim is not None:
|
||
t_dim, h_dim, w_dim = self.fhw_dim
|
||
else:
|
||
# 自动分配的情况
|
||
h_dim = w_dim = 2 * (self.attention_head_dim // 6)
|
||
t_dim = self.attention_head_dim - h_dim - w_dim
|
||
|
||
# 使用正确的split sizes（每个维度的一半）
|
||
freqs = self.freqs.split_with_sizes(
|
||
[
|
||
t_dim // 2, # 时间维度
|
||
h_dim // 2, # 高度维度
|
||
w_dim // 2, # 宽度维度
|
||
],
|
||
dim=1,
|
||
)
|
||
|
||
# 处理causal模式：如果指定了f_end，重新计算ppf和帧范围
|
||
if f_end is not None:
|
||
ppf = f_end - f_start
|
||
frame_slice = slice(f_start, f_end)
|
||
else:
|
||
# 非causal模式：使用从0开始的ppf个帧
|
||
frame_slice = slice(0, ppf)
|
||
|
||
# 步骤2：处理特殊token（如果存在）
|
||
## For other tokens
|
||
if patch_start_idx > 0:
|
||
# 2.1 为特殊token生成位置编码
|
||
# 特殊token位于对角线位置 (f, i, i)，每个特殊token有唯一位置
|
||
# camera: (f, 0, 0), register_0: (f, 1, 1), ..., scale: (f, 5, 5)
|
||
# Shape: (ppf, patch_start_idx, dim)
|
||
freqs_special_f = freqs[0][frame_slice].reshape(ppf, 1, -1).expand(ppf, patch_start_idx, -1) # (ppf, patch_start_idx, dim_f) 帧维度变化
|
||
freqs_special_h = freqs[1][:patch_start_idx].reshape(1, patch_start_idx, -1).expand(ppf, patch_start_idx, -1) # (ppf, patch_start_idx, dim_h) 高度=0,1,2,...
|
||
freqs_special_w = freqs[2][:patch_start_idx].reshape(1, patch_start_idx, -1).expand(ppf, patch_start_idx, -1) # (ppf, patch_start_idx, dim_w) 宽度=0,1,2,...
|
||
freqs_special = torch.cat([freqs_special_f, freqs_special_h, freqs_special_w], dim=-1) # (ppf, patch_start_idx, dim) 拼接三维
|
||
freqs_special = freqs_special.reshape(ppf, patch_start_idx, -1) # (ppf, patch_start_idx, dim)
|
||
|
||
# 2.2 为图像patch生成位置编码
|
||
# Patch位于 (f, patch_start_idx+h, patch_start_idx+w)，h,w 整体偏移 patch_start_idx
|
||
# 这样 patches 与 special tokens 位置不冲突，且 h,w 对称处理
|
||
# Shape: (ppf, pph, ppw, dim)
|
||
freqs_f = freqs[0][frame_slice].reshape(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) # (ppf, pph, ppw, dim_f) 帧维度
|
||
freqs_h = freqs[1][patch_start_idx : patch_start_idx + pph].reshape(1, pph, 1, -1).expand(ppf, pph, ppw, -1) # (ppf, pph, ppw, dim_h) 高度从patch_start_idx开始
|
||
freqs_w = freqs[2][patch_start_idx : patch_start_idx + ppw].reshape(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) # (ppf, pph, ppw, dim_w) 宽度从patch_start_idx开始
|
||
freqs_patches = torch.cat([freqs_f, freqs_h, freqs_w], dim=-1) # (ppf, pph, ppw, dim) 拼接三维
|
||
freqs_patches = freqs_patches.reshape(ppf, pph * ppw, -1) # (ppf, pph * ppw, dim) 展平空间维度
|
||
|
||
# 步骤3：按照正确的顺序组合特殊token和patches
|
||
# 每帧内部顺序：[特殊tokens, patches]
|
||
# Concatenate special tokens and patches for each frame along the second dimension
|
||
# Shape: (ppf, patch_start_idx + pph * ppw, dim)
|
||
freqs = torch.cat([freqs_special, freqs_patches], dim=1) # (ppf, patch_start_idx + pph * ppw, dim)
|
||
|
||
# 步骤4：展平为最终形状并添加batch和head维度
|
||
# Flatten to get final shape: (ppf * (patch_start_idx + pph * ppw), dim)
|
||
freqs = freqs.reshape(ppf * (patch_start_idx + pph * ppw), -1)
|
||
freqs = freqs.unsqueeze(0).unsqueeze(0) # (1, 1, ppf * (patch_start_idx + pph * ppw), dim) 添加batch和head维度
|
||
return freqs
|
||
|
||
# 如果没有特殊token（patch_start_idx == 0），只处理图像patches
|
||
# 所有patches位于 (f, 0:pph, 0:ppw)
|
||
freqs_f = freqs[0][frame_slice].reshape(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) # (ppf, pph, ppw, dim_f) 帧维度
|
||
freqs_h = freqs[1][:pph].reshape(1, pph, 1, -1).expand(ppf, pph, ppw, -1) # (ppf, pph, ppw, dim_h) 高度从0开始
|
||
freqs_w = freqs[2][:ppw].reshape(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) # (ppf, pph, ppw, dim_w) 宽度从0开始
|
||
freqs = torch.cat([freqs_f, freqs_h, freqs_w], dim=-1).reshape(1, 1, ppf * pph * ppw, -1) # (1, 1, ppf * pph * ppw, dim)
|
||
return freqs
|
||
|
||
def apply_rotary_emb(x, freqs):
|
||
"""
|
||
应用旋转位置编码到输入特征
|
||
|
||
核心思想：使用复数乘法实现特征旋转，保持相对位置信息
|
||
|
||
数学原理：
|
||
对于2D向量 [x1, x2]，旋转θ角度可以表示为复数乘法：
|
||
(x1 + ix2) * e^(iθ) = (x1 + ix2) * (cos(θ) + i*sin(θ))
|
||
= (x1*cos(θ) - x2*sin(θ)) + i*(x1*sin(θ) + x2*cos(θ))
|
||
|
||
这等价于旋转矩阵：
|
||
[cos(θ) -sin(θ)] [x1]
|
||
[sin(θ) cos(θ)] [x2]
|
||
|
||
参数：
|
||
- x: 输入特征 [batch, heads, seq_len, head_dim]
|
||
- freqs: 旋转频率（复数） [1, 1, seq_len, head_dim//2]
|
||
|
||
返回：
|
||
- x_out: 旋转后的特征 [batch, heads, seq_len, head_dim]
|
||
|
||
实现步骤：
|
||
1. 将x的每两个连续特征看作一个复数 (real, imag)
|
||
2. 与预计算的复数频率 e^(iθ) 相乘
|
||
3. 转回实数表示
|
||
"""
|
||
# 步骤1：reshape成 [..., head_dim//2, 2] 形式，最后一维表示(real, imag)
|
||
# 例如：[b, h, seq, 64] -> [b, h, seq, 32, 2]
|
||
x_reshaped = x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2)
|
||
|
||
# 步骤2：转换为复数表示 [b, h, seq, 32]
|
||
# 每个元素是 real + imag*i
|
||
x_complex = torch.view_as_complex(x_reshaped)
|
||
|
||
# 步骤3：复数乘法实现旋转
|
||
# x_complex * freqs 相当于将每对特征旋转θ角度
|
||
# freqs已经是 e^(iθ) = cos(θ) + i*sin(θ) 的形式
|
||
x_rotated = x_complex * freqs
|
||
|
||
# 步骤4：转回实数表示 [b, h, seq, 32, 2]
|
||
x_real = torch.view_as_real(x_rotated)
|
||
|
||
# 步骤5：展平最后两维 [b, h, seq, 64]
|
||
x_out = x_real.flatten(3)
|
||
|
||
# 步骤6：转回原始数据类型
|
||
return x_out.to(x.dtype)
|