fuel-nn 0.19.1

Minimalist ML framework.
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

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Lazy-graph Module wrappers.
//!
//! Port of `fuel-nn` over `Tensor`. Each module is a thin wrapper
//! that holds weights and implements `Module::forward`, delegating
//! to the matching `Tensor` primitive. The eager `Module` trait in
//! `fuel-core` is built around `Tensor`; this module mirrors that
//! shape for `Tensor` so downstream lazy ports can build their
//! layer graphs out of named building blocks rather than ad-hoc
//! per-port helpers.

pub mod activation;
pub mod conv;
pub mod embedding;
pub mod init;
pub mod linear;
pub mod lora;
pub mod moe;
pub mod norm;
pub mod quantizable_linear;
pub mod sampling;
pub mod sequential;
pub mod two_proj_attention;

pub use activation::{Elu, Gelu, GeluPytorchTanh, LeakyRelu, Relu, Sigmoid, Silu, Tanh};
pub use conv::{Conv1d, Conv1dConfig, Conv2d, Conv2dConfig};
pub use embedding::Embedding;
pub use linear::{Linear, linear, linear_no_bias};
pub use lora::LoraLinear;
pub use moe::{MoeExpert, MoeLayer, MoeRouter};
pub use norm::{BatchNorm2d, GroupNorm, LayerNorm, RmsNorm};
pub use quantizable_linear::QuantizableLinear;
pub use sequential::Sequential;
pub use two_proj_attention::TwoProjAttention;

use fuel_ir::error::Result;
use fuel_tensor::lazy::Tensor;

/// Single-input `forward` over the lazy-graph tensor. Analogous to
/// the eager `crate::Module` trait, retargeted at `Tensor`.
pub trait Module {
    fn forward(&self, xs: &Tensor) -> Result<Tensor>;
}

impl<F: Fn(&Tensor) -> Result<Tensor>> Module for F {
    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
        self(xs)
    }
}

impl<M: Module> Module for Option<&M> {
    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
        match self {
            None => Ok(xs.clone()),
            Some(m) => m.forward(xs),
        }
    }
}