xn 0.1.1

Another minimalist deep-learning framework optimized for inference.
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

use crate::nn::var_builder::Path;
use crate::{Backend, Result, Tensor, WithDTypeF};

pub struct Linear<T: WithDTypeF, B: Backend> {
    weight: Tensor<T, B>,
    bias: Option<Tensor<T, B>>,
}

impl<T: WithDTypeF, B: Backend> Linear<T, B> {
    pub fn new(weight: Tensor<T, B>) -> Self {
        Self { weight, bias: None }
    }

    pub fn with_bias(self, bias: Tensor<T, B>) -> Self {
        Self { bias: Some(bias), ..self }
    }

    pub fn load<V: std::borrow::Borrow<Path<B>>>(
        vb: V,
        in_features: usize,
        out_features: usize,
    ) -> Result<Self> {
        let vb = vb.borrow();
        let weight = vb.tensor("weight", (out_features, in_features))?;
        Ok(Self::new(weight))
    }

    pub fn load_b<V: std::borrow::Borrow<Path<B>>>(
        vb: V,
        in_features: usize,
        out_features: usize,
    ) -> Result<Self> {
        let vb = vb.borrow();
        let weight = vb.tensor("weight", (out_features, in_features))?;
        let bias = vb.tensor("bias", (out_features,))?;
        Ok(Self::new(weight).with_bias(bias))
    }

    pub fn load_o<V: std::borrow::Borrow<Path<B>>>(
        vb: V,
        in_features: usize,
        out_features: usize,
        bias: bool,
    ) -> Result<Self> {
        let vb = vb.borrow();
        let weight = vb.tensor("weight", (out_features, in_features))?;
        let slf = Self::new(weight);
        if bias {
            let bias = vb.tensor("bias", (out_features,))?;
            Ok(slf.with_bias(bias))
        } else {
            Ok(slf)
        }
    }

    pub fn forward<X: crate::TensorOrView<T, B>>(&self, x: &X) -> Result<Tensor<T, B>> {
        // weight: (out_features, in_features)
        // x: (..., in_features)
        // output: (..., out_features)
        let x = crate::ops::matmul_t(x, &self.weight)?;
        let x = match &self.bias {
            Some(bias) => x.broadcast_add(bias)?,
            None => x,
        };
        Ok(x)
    }

    pub fn device(&self) -> &B {
        self.weight.device()
    }
}