Conv1d for input embedding.
(&self, x: &Tensor)
| 270 | |
| 271 | /// Conv1d for input embedding. |
| 272 | fn embed_conv1d(&self, x: &Tensor) -> Result<Tensor> { |
| 273 | // x: [batch, in_channels, seq] |
| 274 | let (_batch, _in_ch, seq_len) = x.dims3()?; |
| 275 | let pad = self.embed_kernel / 2; |
| 276 | let x_padded = x.pad_with_zeros(2, pad, pad)?; |
| 277 | |
| 278 | // embed_weight: [out_channels, in_channels, kernel_size] |
| 279 | // Standard conv1d (not depthwise) |
| 280 | let mut outputs = Vec::with_capacity(seq_len); |
| 281 | for i in 0..seq_len { |
| 282 | let slice = x_padded.narrow(2, i, self.embed_kernel)?; // [batch, in_ch, kernel] |
| 283 | // For each output position: sum over (in_ch, kernel) dimensions |
| 284 | // slice: [batch, in_ch, kernel], weight: [out_ch, in_ch, kernel] |
| 285 | // output[pos] = einsum('bik,oik->bo', slice, weight) |
| 286 | let slice_expanded = slice.unsqueeze(1)?; // [batch, 1, in_ch, kernel] |
| 287 | let weight_expanded = self.embed_weight.unsqueeze(0)?; // [1, out_ch, in_ch, kernel] |
| 288 | let prod = slice_expanded.broadcast_mul(&weight_expanded)?; // [batch, out_ch, in_ch, kernel] |
| 289 | let summed = prod.sum(candle_core::D::Minus1)?.sum(candle_core::D::Minus1)?; // [batch, out_ch] |
| 290 | outputs.push(summed); |
| 291 | } |
| 292 | let result = Tensor::stack(&outputs, 2)?; // [batch, out_ch, seq] |
| 293 | let bias = self.embed_bias.reshape((1, self.backbone_dim, 1))?; |
| 294 | Ok(result.broadcast_add(&bias)?) |
| 295 | } |
| 296 | |
| 297 | fn layer_norm(&self, x: &Tensor, weight: &Tensor, bias: &Tensor) -> Result<Tensor> { |
| 298 | let mean = x.mean_keepdim(candle_core::D::Minus1)?; |