| 381 | |
| 382 | impl Module for EmbedNd { |
| 383 | fn forward(&self, ids: &Tensor) -> Result<Tensor> { |
| 384 | let n_axes = ids.dim(D::Minus1)?; |
| 385 | let mut emb = Vec::with_capacity(n_axes); |
| 386 | for idx in 0..n_axes { |
| 387 | let pos = ids.get_on_dim(D::Minus1, idx)?; |
| 388 | let dim = self.axes_dim[idx]; |
| 389 | let half = dim / 2; |
| 390 | let dev = pos.device(); |
| 391 | |
| 392 | let inv_freq = Tensor::from_slice(&self.inv_freqs[idx], (1, 1, half), dev)?; |
| 393 | let inv_freq = inv_freq.to_dtype(pos.dtype())?; |
| 394 | let freqs = pos.unsqueeze(2)?.broadcast_mul(&inv_freq)?; |
| 395 | let cos = freqs.cos()?; |
| 396 | let sin = freqs.sin()?; |
| 397 | let out = Tensor::stack(&[&cos, &sin.neg()?, &sin, &cos], 3)?; |
| 398 | let (b, n, d, _ij) = out.dims4()?; |
| 399 | let r = out.reshape((b, n, d, 2, 2))?; |
| 400 | emb.push(r); |
| 401 | } |
| 402 | let emb = Tensor::cat(&emb, 2)?; |
| 403 | emb.unsqueeze(1) |
| 404 | } |
| 405 | } |
| 406 | |
| 407 | // ── MLP Embedder ─────────────────────────────────────────────────────────────── |