microflow/ops/
softmax.rs

1use crate::activation;
2use crate::quantize::Quantized;
3use crate::tensor::Tensor2D;
4use libm::expf;
5use simba::scalar::SupersetOf;
6
7/// Performs the Softmax activation function as an operator.
8/// Returns a 2-dimensional output tensor containing the result of the operation.
9///
10/// # Arguments
11/// * `input` - The 2-dimensional input tensor
12/// * `output_scale` - The scale of the resulting output tensor
13/// * `output_zero_point` - The zero point of the resulting output tensor
14///
15pub fn softmax<T: Quantized, const ROWS: usize, const COLS: usize>(
16    input: Tensor2D<T, ROWS, COLS, 1>,
17    output_scale: [f32; 1],
18    output_zero_point: [T; 1],
19) -> Tensor2D<T, ROWS, COLS, 1> {
20    let exp = input.buffer.map(|e| f32::from_subset(&e) * input.scale[0]);
21    let sum = exp.map(expf).sum();
22    Tensor2D::new(
23        exp.map(|e| activation::softmax(e, sum, output_scale[0], output_zero_point[0])),
24        output_scale,
25        output_zero_point,
26    )
27}
28
29#[cfg(test)]
30mod tests {
31    use super::*;
32    use nalgebra::matrix;
33
34    const INPUT: Tensor2D<i8, 2, 3, 1> = Tensor2D {
35        buffer: matrix![
36            1, 2, 3;
37            4, 5, 6
38        ],
39        scale: [0.7],
40        zero_point: [8],
41    };
42    const OUTPUT_SCALE: [f32; 1] = [0.9];
43    const OUTPUT_ZERO_POINT: [i8; 1] = [10];
44    const OUTPUT: Tensor2D<i8, 2, 3, 1> = Tensor2D {
45        buffer: matrix![
46            10, 10, 10;
47            10, 10, 11
48        ],
49        scale: OUTPUT_SCALE,
50        zero_point: OUTPUT_ZERO_POINT,
51    };
52
53    #[test]
54    fn softmax_layer() {
55        assert_eq!(softmax(INPUT, OUTPUT_SCALE, OUTPUT_ZERO_POINT), OUTPUT);
56    }
57}
microflow/ops/softmax.rs

microflow/ops/
softmax.rs
