refactor: 拆分 module.py 为 components 子包

- rope/linear/norm/embedding/mlp/attention/decoder_block 各自独立文件
- 依赖单向无循环
- 公开接口不变,外部无需修改
This commit is contained in:
2026-05-15 20:08:36 +08:00
parent 19532440b4
commit ef25efffa2
10 changed files with 205 additions and 154 deletions
+18
View File
@@ -0,0 +1,18 @@
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from astrai.model.components.linear import Linear
class MLP(nn.Module):
def __init__(self, dim: int, dim_feed_forward: int):
super().__init__()
self.up = Linear(dim, dim_feed_forward)
self.gate = Linear(dim, dim_feed_forward)
self.down = Linear(dim_feed_forward, dim)
def forward(self, x: Tensor) -> Tensor:
gated = self.up(x) * F.silu(self.gate(x))
out = self.down(gated)
return out