| 27 | } |
| 28 | |
| 29 | pub fn forward(&self, t: &Tensor) -> Result<Tensor> { |
| 30 | // Sinusoidal embedding: t → (batch, 256) |
| 31 | let half_dim = 128; |
| 32 | let emb = { |
| 33 | // Compute frequency vector on CPU — avoids arange + to_dtype + mul + exp tensor ops |
| 34 | let decay = -f64::ln(10000.0) / half_dim as f64; |
| 35 | let freq_data: Vec<f32> = (0..half_dim).map(|j| (j as f64 * decay).exp() as f32).collect(); |
| 36 | let freq = Tensor::new(freq_data.as_slice(), t.device())?.unsqueeze(0)?; |
| 37 | let t_f32 = t.to_dtype(DType::F32)?; |
| 38 | let args = t_f32.unsqueeze(1)?.broadcast_mul(&freq)?; |
| 39 | Tensor::cat(&[args.cos()?, args.sin()?], D::Minus1)?.to_dtype(t.dtype())? |
| 40 | }; |
| 41 | // MLP: 256 → hidden → hidden with SiLU |
| 42 | let h = self.backend.linear_forward(&emb, &self.mlp_0_weight, None)?; |
| 43 | let h = self.backend.silu(&h)?; |
| 44 | self.backend.linear_forward(&h, &self.mlp_2_weight, None) |
| 45 | } |
| 46 | } |
| 47 | |
| 48 | /// SwiGLU feed-forward network. |