hpt_traits/ops/
normalization.rs

1use hpt_common::error::base::TensorError;
2use hpt_common::shape::shape::Shape;
3use hpt_types::into_scalar::Cast;
4
5/// A trait contains normalization operations
6pub trait NormalizationOps {
7    /// The type of the output tensor
8    type Output;
9    /// The type of the output meta
10    type OutputMeta;
11
12    /// Applies Layer Normalization over a specified axes.
13    ///
14    /// ## Parameters:
15    /// `normalized_shape`: shape that must match the dimension size from input tensor shape (from right to left)
16    ///
17    /// `gamma`: Optional scale tensor of shape `[normalized_shape]`
18    ///
19    /// `beta`: Optional bias tensor of shape `[normalized_shape]`
20    ///
21    /// `eps`: A value added to the denominator for numerical stability.
22    ///
23    /// ## Example:
24    /// ```rust
25    /// let x = Tensor::<f32>::randn(&[2, 3, 4])?;
26    /// let gamma = Tensor::<f32>::ones(&[4])?;
27    /// let beta = Tensor::<f32>::zeros(&[4])?;
28    /// let result = x.layernorm(&[4], Some(&gamma), Some(&beta), 1e-5)?;
29    /// ```
30    fn layernorm<S: Into<Shape>>(
31        &self,
32        normalized_shape: S,
33        gamma: Option<&Self::Output>,
34        beta: Option<&Self::Output>,
35        eps: Self::OutputMeta,
36    ) -> Result<Self::Output, TensorError>
37    where
38        usize: Cast<Self::OutputMeta>;
39
40    /// Applies the softmax function to the input tensor along the specified dimension.
41    /// The softmax function normalizes the input to a probability distribution, such that each element is in the range [0, 1] and all elements sum to 1.
42    ///
43    /// ## Parameters:
44    /// `dim`: The dimension along which to apply the softmax.
45    ///
46    /// ## Example:
47    /// ```rust
48    /// let x = Tensor::<f32>::new(&[[-1.0, 0.0, 1.0], [2.0, 3.0, 4.0]]);
49    /// let result = x.softmax(1)?;
50    /// ```
51    fn softmax(&self, axis: i64) -> Result<Self::Output, TensorError>;
52
53    /// Applies the log-softmax function to the input tensor along the specified dimension.
54    /// The log-softmax function is equivalent to applying the logarithm to the output of the softmax function, but is more numerically stable when computed directly.
55    ///
56    /// ## Parameters:
57    /// `dim`: The dimension along which to apply the log-softmax.
58    ///
59    /// ## Example:
60    /// ```rust
61    /// let x = Tensor::<f32>::new(&[[-1.0, 0.0, 1.0], [2.0, 3.0, 4.0]]);
62    /// let result = x.log_softmax(1)?;
63    /// ```
64    fn log_softmax(&self, axis: i64) -> Result<Self::Output, TensorError>;
65}
hpt_traits/ops/normalization.rs

hpt_traits/ops/
normalization.rs
