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

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

/// A representation of a temporal (time series) layer builder.
pub struct Temporal<const RANK: usize, const SIZE: usize> {
    /// The shape of the layer.
    pub shape: [usize; RANK],

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

impl<const RANK: usize, const SIZE: usize> Temporal<RANK, SIZE> {
    /// Builds the structure into a compute layer.
    #[must_use]
    pub fn build(self) -> TemporalLayer<RANK, SIZE> {
        TemporalLayer {
            attn: Tensor::new(&self.shape, &self.data).with_grad(),
        }
    }
}

/// An expanded representation of a temporal (time series) layer.
///
/// Implements scaled dot-product self-attention mechanism.
/// The weights are organized as: [`Q | K | V | Output`] where each is `d_model x d_model`.
/// `PROJ_SIZE` must equal `SIZE / 4` (each projection is `PROJ_SIZE` elements).
pub struct TemporalLayer<const RANK: usize, const SIZE: usize> {
    attn: WithGrad<Tensor<RANK, SIZE>>,
}

impl<const RANK: usize, const SIZE: usize> TemporalLayer<RANK, SIZE> {
    /// Forwards the temporal layer.
    #[must_use]
    #[cfg(feature = "dyntensor")]
    pub fn forward<'a>(
        &'a self,
        input: &'a WithGrad<Tensor<0, 0>>,
    ) -> (Tensor<0, 0>, Backward<'a, 0, 0, 0, 0>) {
        let scale = 1.0 / libm::sqrtf(self.attn.get_value().shape()[RANK - 1] as f32) as TensorFloat;
        let scaled = self.attn.get_value().map(|&x| x * scale).with_grad();

        let (attn_sm, _) = softmax(&scaled);

        let (out, _) = matmul(input, &attn_sm.clone().with_grad());

        let back = move |grad_output: Tensor<0, 0>| {
            let attn_t = attn_sm.transpose();
            grad_output.matmul(&attn_t)
        };

        (out, Backward::Unary(alloc::boxed::Box::new(back)))
    }

    /// Forwards the temporal layer.
    #[must_use]
    #[cfg(not(feature = "dyntensor"))]
    pub fn forward<'a, const IN_SIZE: usize, const OUT_SIZE: usize>(
        &'a self,
        input: &'a WithGrad<Tensor<RANK, IN_SIZE>>,
    ) -> (Tensor<RANK, OUT_SIZE>, Backward<'a, RANK, SIZE, OUT_SIZE, IN_SIZE>) {
        let scale = 1.0 / libm::sqrtf(self.attn.get_value().shape()[RANK - 1] as f32) as TensorFloat;
        let scaled = self.attn.get_value().map(|x| x * scale).with_grad();

        let (attn_sm, _) = softmax::<SIZE, SIZE, RANK>(&scaled);

        let (out, _) = matmul(input, &attn_sm.clone().with_grad());

        let back = move |grad_output: Tensor<RANK, OUT_SIZE>| {
            let attn_t = attn_sm.transpose();
            grad_output.matmul(&attn_t)
        };

        #[cfg(feature = "alloc")]
        {
            (out, Backward::Unary(alloc::boxed::Box::new(back)))
        }
        #[cfg(not(feature = "alloc"))]
        {
            (out, Backward::Unary(box_closure::OpaqueFn::new(back)))
        }
    }

    /// Differentiates the temporal layer.
    #[must_use]
    #[cfg(feature = "dyntensor")]
    pub fn backward(
        &self,
        grad_output: Tensor<0, 0>,
        back: Backward<'_, 0, 0, 0, 0>,
    ) -> (Tensor<0, 0>, Option<Tensor<0, 0>>) {
        let d_input = match back {
            Backward::Binary(_) => unreachable!("Temporal never has a binary closure"),
            Backward::Unary(f) => f.invoke(grad_output),
        };
        (d_input, None)
    }

    /// Differentiates the temporal layer.
    #[must_use]
    #[cfg(not(feature = "dyntensor"))]
    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, SIZE>>) {
        let d_input = match back {
            Backward::Binary(_) => unreachable!("Temporal never has a binary closure"),
            Backward::Unary(f) => f.invoke(grad_output),
        };
        (d_input, None)
    }
}

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

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