bunsen 0.29.1

bunsen is a batteries included common library for burn
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

use burn::{
    Tensor,
    nn::{
        LayerNorm,
        attention::{
            MhaInput,
            MhaOutput,
            MultiHeadAttention,
        },
    },
    prelude::{
        Backend,
        Bool,
    },
};

/// Computes layer normalized self-attn.
///
/// # Arguments
/// * `layer_norm` - `LayerNorm`.
/// * `mh_attn` - `MultiHeadAttention`.
/// * `x` - `[batch, seq_len, d_model]` input.
/// * `mask` - Optional `[batch, seq_len, seq_len]` attention mask.
///
/// # Returns
/// `RdabForwardRecord` - forward record.
/// * `fr.output` : `[batch, seq_len, d_model]`.
/// * `fr.ca_weights` : `[batch, n_heads, seq_len, seq_len]`.
pub fn layer_norm_self_attn<B: Backend>(
    layer_norm: &LayerNorm<B>,
    mh_attn: &MultiHeadAttention<B>,
    x: Tensor<B, 3>,
    mask: Option<Tensor<B, 3, Bool>>,
) -> MhaOutput<B> {
    #[cfg(any(debug_assertions, test))]
    {
        use crate::contracts::*;
        let d_model = mh_attn.d_model;
        assert_eq!(
            d_model,
            layer_norm.gamma.dims()[0],
            "layer_norm dims ({}) != d_model ({d_model})",
            layer_norm.gamma.dims()[0],
        );

        let [batch, seq_len] = unpack_shape_contract!(
            ["batch", "seq_len", "d_model"],
            &x,
            &["batch", "seq_len"],
            &[("d_model", d_model)]
        );

        if let Some(mask) = &mask {
            let [mask_batch] = unpack_shape_contract!(
                ["mask_batch", "seq_len", "seq_len"],
                mask,
                &["mask_batch"],
                &[("seq_len", seq_len)]
            );
            if mask_batch != 1 {
                assert_eq!(
                    mask_batch, batch,
                    "batch sizes not broadcastable {batch} vs {mask_batch}"
                );
            }

            assert_shape_contract!(["b", "seq_len", "seq_len"], mask, &[("seq_len", seq_len)]);
        }
    }

    let input = MhaInput::self_attn(layer_norm.forward(x));
    let input = match mask {
        Some(mask) => input.mask_attn(mask),
        None => input,
    };
    mh_attn.forward(input)
}

/// Computes layer normalized cross-attn.
///
/// # Arguments
/// * `layer_norm` - `LayerNorm`.
/// * `mh_attn` - `MultiHeadAttention`.
/// * `x` - `[batch, seq_len, d_model]` input.
/// * `xa` - `[batch, seq_len, d_model]` cross-attention input.
///
/// # Returns
/// `RdabForwardRecord` - forward record.
/// * `fr.output` : `[batch, seq_len, d_model]`.
/// * `fr.ca_weights` : `[batch, n_heads, seq_len, seq_len]`.
pub fn layer_norm_cross_attn<B: Backend>(
    layer_norm: &LayerNorm<B>,
    mh_attn: &MultiHeadAttention<B>,
    x: Tensor<B, 3>,
    xa: Tensor<B, 3>,
) -> MhaOutput<B> {
    #[cfg(any(debug_assertions, test))]
    {
        crate::contracts::define_shape_contract!(CONTRACT, ["batch", "seq_len", "d_model"]);
        let d_model = mh_attn.d_model;
        assert_eq!(
            d_model,
            layer_norm.gamma.dims()[0],
            "layer_norm dims ({}) != d_model ({d_model})",
            layer_norm.gamma.dims()[0],
        );

        let [batch, seq_len] =
            CONTRACT.unpack_shape(&x, &["batch", "seq_len"], &[("d_model", d_model)]);

        CONTRACT.assert_shape(
            &xa,
            &[("batch", batch), ("seq_len", seq_len), ("d_model", d_model)],
        );
    };

    mh_attn.forward(MhaInput::new(layer_norm.forward(x.clone()), xa.clone(), xa))
}