| 55 | """ |
| 56 | |
| 57 | def __init__( |
| 58 | self, |
| 59 | in_features: int, |
| 60 | hidden_features: Optional[int] = None, |
| 61 | out_features: Optional[int] = None, |
| 62 | act_layer: Optional[Callable[..., nn.Module]] = None, |
| 63 | drop: float = 0.0, |
| 64 | bias: bool = True, |
| 65 | align_to: int = 8, |
| 66 | device=None, |
| 67 | ) -> None: |
| 68 | super().__init__() |
| 69 | out_features = out_features or in_features |
| 70 | hidden_features = hidden_features or in_features |
| 71 | d = int(hidden_features * 2 / 3) |
| 72 | swiglu_hidden_features = d + (-d % align_to) |
| 73 | self.w1 = nn.Linear(in_features, swiglu_hidden_features, bias=bias, device=device) |
| 74 | self.w2 = nn.Linear(in_features, swiglu_hidden_features, bias=bias, device=device) |
| 75 | self.w3 = nn.Linear(swiglu_hidden_features, out_features, bias=bias, device=device) |
| 76 | |
| 77 | def forward(self, x: Tensor) -> Tensor: |
| 78 | x1 = self.w1(x) |