| 137 | } |
| 138 | |
| 139 | fn forward(&self, x: &Tensor) -> Result<Tensor> { |
| 140 | let channels = x.dim(1)?; |
| 141 | |
| 142 | let h = self.backend.rms_norm_channel(x, &self.norm_weight, self.eps)?; |
| 143 | let zeros = Tensor::zeros((h.dim(0)?, channels, 6), h.dtype(), h.device())?; |
| 144 | let h = self.backend.depthwise_conv1d_bias_ctx( |
| 145 | &zeros, &h, &self.mixer_weight, &self.mixer_bias, 7, channels, |
| 146 | )?; |
| 147 | let x = self.backend.add_scaled(x, &h, &self.gamma)?; |
| 148 | |
| 149 | let h = self.backend.rms_norm_channel(&x, &self.ffn_norm_weight, self.eps)?; |
| 150 | let h = h.transpose(1, 2)?; |
| 151 | let h = self.backend.linear_forward(&h, &self.ffn_linear1_weight, self.ffn_linear1_bias.as_ref())?; |
| 152 | let h = self.backend.gelu(&h)?; |
| 153 | let h = self.backend.linear_forward(&h, &self.ffn_linear2_weight, self.ffn_linear2_bias.as_ref())?; |
| 154 | let h = h.transpose(1, 2)?; |
| 155 | self.backend.add_scaled(&x, &h, &self.ffn_gamma) |
| 156 | } |
| 157 | |
| 158 | /// Forward with streaming cache: uses cached context instead of zero-padding. |
| 159 | fn forward_cached(&self, x: &Tensor, cache: &mut StreamingConvCache) -> Result<Tensor> { |