bunsen 0.27.0

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

use burn::{
    Tensor,
    nn::Embedding,
    prelude::Backend,
};

/// Argument to [unembed].
pub trait EmbeddingArg<B: Backend> {
    /// The embedding layer's weight.
    fn weight(&self) -> Tensor<B, 2>;
}

impl<B: Backend> EmbeddingArg<B> for &Embedding<B> {
    fn weight(&self) -> Tensor<B, 2> {
        self.weight.val()
    }
}

impl<B: Backend> EmbeddingArg<B> for Tensor<B, 2> {
    fn weight(&self) -> Tensor<B, 2> {
        self.clone()
    }
}

/// Inverts an Embedding layer to get logits.
///
/// # Arguments
/// * `emb` - The `[n_vocab, d_model]`, embedding layer to invert.
/// * `x` - The `[batch, seq_len, d_model]` tensor to unembed.
///
/// # Returns
/// The `[batch, seq_len, n_vocab]` logits tensor.
pub fn unembed<B: Backend, A: EmbeddingArg<B>>(
    emb: A,
    x: Tensor<B, 3>,
) -> Tensor<B, 3> {
    x.matmul(emb.weight().transpose().unsqueeze::<3>())
}

#[cfg(test)]
mod tests {
    use burn::tensor::{
        Distribution,
        Int,
    };
    use serial_test::serial;

    use super::*;
    use crate::{
        ops::embedding::identity_embedding,
        support::testing::PerformanceBackend,
    };

    #[test]
    #[serial]
    fn test_embedding_inverse_to_logits() {
        type B = PerformanceBackend;
        let device = Default::default();

        let n_embedding = 10;

        let embedding = identity_embedding(n_embedding, &device);

        let batch = 2;
        let seq_len = 20;

        let x: Tensor<B, 2, Int> = Tensor::random(
            [batch, seq_len],
            Distribution::Uniform(0.0, n_embedding as f64),
            &device,
        );

        let e = embedding.forward(x.clone());

        let logits = unembed(&embedding, e.clone());
        let logits2 = unembed(embedding.weight.val(), e.clone());
        logits
            .clone()
            .into_data()
            .assert_eq(&logits2.into_data(), true);

        let y: Tensor<B, 2, Int> = logits.argmax(2).squeeze_dim(2);

        y.into_data().assert_eq(&x.into_data(), true);
    }
}