mirror of
https://github.com/huggingface/candle.git
synced 2025-06-20 12:06:35 +00:00
45 lines
1.5 KiB
Rust
45 lines
1.5 KiB
Rust
use candle::{Result, Tensor};
|
|
|
|
/// Applies the softmax function to the input tensor, rescaling the element so that elements on
|
|
/// a slice of fixed index on dimension `dim` are between 0 and 1 and sum to 1.
|
|
///
|
|
/// ```rust
|
|
/// use candle::{Tensor, Device};
|
|
/// let a = Tensor::new(&[[0f32, 1., 0., 1.], [-2., 2., 3., -3.]], &Device::Cpu)?;
|
|
/// let a = candle_nn::ops::softmax(&a, 1)?;
|
|
/// assert_eq!(
|
|
/// a.to_vec2::<f32>()?,
|
|
/// &[
|
|
/// [0.13447072, 0.3655293, 0.13447072, 0.3655293],
|
|
/// [0.0048928666, 0.26714146, 0.7261658, 0.0017999851]
|
|
/// ]);
|
|
/// # Ok::<(), candle::Error>(())
|
|
/// ```
|
|
pub fn softmax<D: candle::shape::Dim>(xs: &Tensor, dim: D) -> Result<Tensor> {
|
|
let dim = dim.to_index(xs.shape(), "softmax")?;
|
|
let max = xs.max_keepdim(dim)?;
|
|
let diff = xs.broadcast_sub(&max)?;
|
|
let num = diff.exp()?;
|
|
let den = num.sum_keepdim(dim)?;
|
|
num.broadcast_div(&den)
|
|
}
|
|
|
|
pub fn log_softmax<D: candle::shape::Dim>(xs: &Tensor, d: D) -> Result<Tensor> {
|
|
let d = d.to_index(xs.shape(), "log-softmax")?;
|
|
let max = xs.max_keepdim(d)?;
|
|
let diff = xs.broadcast_sub(&max)?;
|
|
let sum_exp = diff.exp()?.sum_keepdim(d)?;
|
|
let log_sm = diff.broadcast_sub(&sum_exp.log()?)?;
|
|
Ok(log_sm)
|
|
}
|
|
|
|
pub fn silu(xs: &Tensor) -> Result<Tensor> {
|
|
// TODO: Should we have a specialized op for this?
|
|
xs / (xs.neg()?.exp()? + 1.0)?
|
|
}
|
|
|
|
pub fn sigmoid(xs: &Tensor) -> Result<Tensor> {
|
|
// TODO: Should we have a specialized op for this?
|
|
(xs.neg()?.exp()? + 1.0)?.recip()
|
|
}
|