Method
__init__
(
self,
dim: int,
inner_dim: int,
multiple_of: int | None = 256,
ffn_dim_multiplier: float | None = None,
)
Source from the content-addressed store, hash-verified
| 1096 | """ |
| 1097 | |
| 1098 | def __init__( |
| 1099 | self, |
| 1100 | dim: int, |
| 1101 | inner_dim: int, |
| 1102 | multiple_of: int | None = 256, |
| 1103 | ffn_dim_multiplier: float | None = None, |
| 1104 | ): |
| 1105 | super().__init__() |
| 1106 | # custom hidden_size factor multiplier |
| 1107 | if ffn_dim_multiplier is not None: |
| 1108 | inner_dim = int(ffn_dim_multiplier * inner_dim) |
| 1109 | inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of) |
| 1110 | |
| 1111 | self.linear_1 = nn.Linear( |
| 1112 | dim, |
| 1113 | inner_dim, |
| 1114 | bias=False, |
| 1115 | ) |
| 1116 | self.linear_2 = nn.Linear( |
| 1117 | inner_dim, |
| 1118 | dim, |
| 1119 | bias=False, |
| 1120 | ) |
| 1121 | self.linear_3 = nn.Linear( |
| 1122 | dim, |
| 1123 | inner_dim, |
| 1124 | bias=False, |
| 1125 | ) |
| 1126 | self.silu = FP32SiLU() |
| 1127 | |
| 1128 | def forward(self, x): |
| 1129 | return self.linear_2(self.silu(self.linear_1(x)) * self.linear_3(x)) |
Callers
nothing calls this directly
Tested by
no test coverage detected