briny_ai 0.5.0

A tiny & efficient AI inference engine
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

use super::{
    softmax, Backward, Closure, Layer, Tensor, TensorFloat, WithGrad,
};

/// Defines a layer that uses softmax to propagate probability distributions.
pub struct Softmax<const RANK: usize, const SIZE: usize> {
    /// The shape of the layer.
    pub shape: [usize; RANK],

    /// The data of the layer.
    pub data: [TensorFloat; SIZE],

    /// The temperature of the softmax layer.
    pub temp: f32,
}

impl<const RANK: usize, const SIZE: usize>
    Softmax<RANK, SIZE>
{
    /// Builds the structure into a compute layer.
    #[must_use]
    pub const fn build(self) -> SoftmaxLayer<RANK, SIZE> {
        SoftmaxLayer {
            temp: self.temp,
        }
    }
}

/// A layer that uses softmax to propagate probability distributions.
pub struct SoftmaxLayer<const RANK: usize, const SIZE: usize> {
    temp: f32,
}

impl<const RANK: usize, const SIZE: usize> SoftmaxLayer<RANK, SIZE> {
    /// Gets the temperature of the layer.
    #[must_use]
    pub const fn get_temp_relative(&self, _: f32) -> f32 {
        self.temp
    }

    /// Updates the temperature of the layer.
    pub fn update_temp<F>(&mut self, f: F)
    where 
        F: Fn(f32) -> f32,
    {
        self.temp = f(self.temp);
    }
}

#[cfg(feature = "dyntensor")]
impl<const RANK: usize, const SIZE: usize> SoftmaxLayer<RANK, SIZE> {
    /// Forwards the softmax layer.
    #[inline]
    #[must_use]
    pub fn forward<'a>(
        &'a self,
        input: &'a WithGrad<Tensor<RANK, 0>>,
    ) -> (Tensor<RANK, 0>, Backward<'a, RANK, 0, 0, 0>) {
        let (out, f) = softmax(input);

        (out, Backward::Unary(f))
    }

    /// Differentiates the softmax layer with the provided closure.
    #[inline]
    #[must_use]
    #[allow(clippy::unnecessary_wraps)]
    pub fn backward(
        &self,
        grad_output: Tensor<RANK, 0>,
        back: Backward<'_, RANK, 0, 0, 0>,
    ) -> (Tensor<RANK, 0>, Option<Tensor<RANK, 0>>) {
        // call the backward closure and return the results
        match back {
            Backward::Unary(f) => {
                let grad_in = f.invoke(grad_output);
                (grad_in, None)
            }
            Backward::Binary(_) => {
                unreachable!("Softmax never has a binary closure");
            }
        }
    }
}

#[cfg(not(feature = "dyntensor"))]
impl<const RANK: usize, const SIZE: usize> SoftmaxLayer<RANK, SIZE> {
    /// Forwards the softmax layer.
    #[inline]
    #[must_use]
    pub fn forward<'a, const OUT_SIZE: usize>(
        &'a self,
        input: &'a WithGrad<Tensor<RANK, SIZE>>,
    ) -> (Tensor<RANK, OUT_SIZE>, Backward<'a, RANK, OUT_SIZE, OUT_SIZE, SIZE>) {
        let (out, f) = softmax(input);

        (out, Backward::Unary(f))
    }

    /// Differentiates the softmax layer with the provided closure.
    #[inline]
    #[must_use]
    #[allow(clippy::unnecessary_wraps)]
    pub fn backward<const IN_SIZE: usize, const OUT_SIZE: usize>(
        &self,
        grad_output: Tensor<RANK, IN_SIZE>,
        back: Backward<'_, RANK, SIZE, IN_SIZE, OUT_SIZE>,
    ) -> (Tensor<RANK, OUT_SIZE>, Option<Tensor<RANK, OUT_SIZE>>) {
        // call the backward closure and return the results
        match back {
            Backward::Unary(f) => {
                let grad_in = f.invoke(grad_output);
                (grad_in, None)
            }
            Backward::Binary(_) => {
                unreachable!("Softmax never has a binary closure");
            }
        }
    }
}

impl<const RANK: usize, const IN_SIZE: usize> Layer<RANK, IN_SIZE, 0>
    for SoftmaxLayer<RANK, IN_SIZE>
{
    #[inline]
    fn weights(&self) -> [&WithGrad<Tensor<RANK, IN_SIZE>>; 0] {
        []
    }

    #[inline]
    fn weights_mut(&mut self) -> [&mut WithGrad<Tensor<RANK, IN_SIZE>>; 0] {
        []
    }
}