fastnn 0.3.3

A robust deep learning library in Rust with CUDA kernels
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

//! Embedding table and positional encoding.

use crate::tensor::{Tensor, ops::index::index_from_f32};

use super::module::Module;
use super::param::Param;

const INIT_STD: f32 = 0.02;
/// Base of the geometric wavelength progression in the sinusoidal encoding.
const WAVELENGTH_BASE: f32 = 10000.0;

/// A learnable `[vocab, dim]` table looked up by token id.
pub struct Embedding {
    pub weight: Param,
    vocab: usize,
    dim: usize,
}

impl Embedding {
    pub fn new(vocab: usize, dim: usize) -> Embedding {
        // small init so embeddings don't dominate the residual stream early on
        Embedding {
            weight: Param::new(Tensor::randn(&[vocab, dim]).mul_scalar(INIT_STD)),
            vocab,
            dim,
        }
    }

    /// Look up a flat list of ids, giving `[ids.len(), dim]`.
    pub fn lookup(&self, ids: &[usize]) -> Tensor {
        self.weight.tensor().index_select(ids)
    }

    /// Look up a `[batch, sequence]` id grid, giving `[batch, sequence, dim]`.
    pub fn lookup_batch(&self, ids: &[Vec<usize>]) -> Tensor {
        let (batch, sequence) = (ids.len(), ids.first().map_or(0, |row| row.len()));
        let flat: Vec<usize> = ids
            .iter()
            .inspect(|row| assert_eq!(row.len(), sequence, "lookup_batch: ragged id grid"))
            .flatten()
            .copied()
            .collect();
        self.lookup(&flat)
            .reshape(&[batch as i64, sequence as i64, self.dim as i64])
    }

    pub fn vocab(&self) -> usize {
        self.vocab
    }

    pub fn dim(&self) -> usize {
        self.dim
    }
}

impl Module for Embedding {
    /// Treats finite, non-negative integer values as ids and panics on anything
    /// else. Prefer [`lookup`](Embedding::lookup).
    fn forward(&self, input: &Tensor) -> Tensor {
        let ids: Vec<usize> = input.to_vec().into_iter().map(index_from_f32).collect();
        let mut shape: Vec<i64> = input.shape().iter().map(|&d| d as i64).collect();
        shape.push(self.dim as i64);
        self.lookup(&ids).reshape(&shape)
    }

    fn named_parameters(&self) -> Vec<(String, Param)> {
        vec![("weight".into(), self.weight.clone())]
    }
}

/// Fixed sinusoidal position signal, added to token embeddings.
pub struct PositionalEncoding {
    table: Tensor,
    dim: usize,
}

impl PositionalEncoding {
    pub fn new(dim: usize, max_len: usize) -> PositionalEncoding {
        let mut table = vec![0.0f32; max_len * dim];
        for position in 0..max_len {
            for i in (0..dim).step_by(2) {
                let angle = position as f32 / WAVELENGTH_BASE.powf(i as f32 / dim as f32);
                table[position * dim + i] = angle.sin();
                if i + 1 < dim {
                    table[position * dim + i + 1] = angle.cos();
                }
            }
        }
        PositionalEncoding {
            table: Tensor::from_vec(table, &[max_len, dim]),
            dim,
        }
    }
}

impl Module for PositionalEncoding {
    /// Adds positions to a `[batch, sequence, dim]` input.
    fn forward(&self, input: &Tensor) -> Tensor {
        let sequence = input.dim(1);
        assert!(
            sequence <= self.table.dim(0),
            "PositionalEncoding: sequence {sequence} exceeds max_len {}",
            self.table.dim(0)
        );
        let positions: Vec<usize> = (0..sequence).collect();
        let slice = self
            .table
            .to(input.device())
            .index_select(&positions)
            .reshape(&[1, sequence as i64, self.dim as i64]);
        input.add(&slice)
    }
}