Softmax over the given dimension. Uses raw f32 computation for last-dim F32 (avoids CustomOp dispatch), fused kernel for other dtypes, generic path for non-last dim.
(&self, x: &Tensor, dim: usize)
| 405 | /// Uses raw f32 computation for last-dim F32 (avoids CustomOp dispatch), |
| 406 | /// fused kernel for other dtypes, generic path for non-last dim. |
| 407 | fn softmax(&self, x: &Tensor, dim: usize) -> Result<Tensor> { |
| 408 | if dim == x.rank() - 1 { |
| 409 | // Fast path: F32 last-dim softmax on raw data |
| 410 | if x.dtype() == DType::F32 { |
| 411 | let shape = x.dims(); |
| 412 | let last = *shape.last().unwrap_or(&0); |
| 413 | let data = x.contiguous()?.flatten_all()?.to_vec1::<f32>()?; |
| 414 | let rows = data.len() / last; |
| 415 | let mut out = vec![0f32; data.len()]; |
| 416 | for r in 0..rows { |
| 417 | let off = r * last; |
| 418 | let row = &data[off..off + last]; |
| 419 | let max = row.iter().copied().fold(f32::NEG_INFINITY, f32::max); |
| 420 | let mut sum = 0f32; |
| 421 | for i in 0..last { |
| 422 | let e = (row[i] - max).exp(); |
| 423 | out[off + i] = e; |
| 424 | sum += e; |
| 425 | } |
| 426 | let inv_sum = 1.0 / sum; |
| 427 | for i in 0..last { |
| 428 | out[off + i] *= inv_sum; |
| 429 | } |
| 430 | } |
| 431 | return Tensor::from_vec(out, shape, x.device()); |
| 432 | } |
| 433 | candle_nn::ops::softmax_last_dim(x) |
| 434 | } else { |
| 435 | let max = x.max_keepdim(dim)?; |
| 436 | let exp = x.broadcast_sub(&max)?.exp()?; |
| 437 | let sum = exp.sum_keepdim(dim)?; |
| 438 | exp.broadcast_div(&sum) |
| 439 | } |
| 440 | } |
| 441 | |
| 442 | /// Rotary position embedding: apply cos/sin rotation to tensor. |
| 443 | /// `cos` and `sin` have shape `(seq_len, head_dim/2)` or compatible broadcast shape. |