lumen-core 0.6.1

A tiny ML framework
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

mod construct;
mod indexer;
mod iter;
pub mod display;
mod shape;
mod arith;
mod matmul;
mod reduce;
mod broadcast;
mod convert;
mod boolean;
mod fused;
mod nn;
mod variable;

pub use construct::ToTensor;
pub use indexer::{Slice, IndexOp};
pub use iter::*;
pub use indexer::*;
pub use variable::*;

use crate::{FloatDType, IntTensor, Tensor, WithDType};

#[derive(Clone)]
pub enum Op<T: FloatDType> {
    Binary(Tensor<T>, Tensor<T>, BinaryOp),
    BinaryScalarRhs(Tensor<T>, T, BinaryOp),
    BinaryScalarLhs(T, Tensor<T>, BinaryOp),
    Unary(Tensor<T>, UnaryOp<T>),
    Reduce(Tensor<T>, ReduceOp, Vec<usize>),
    Matmul(Tensor<T>, Tensor<T>),
    Broadcast(Tensor<T>),
    Narrow(Tensor<T>, usize, usize, usize),
    Slice(Tensor<T>, usize, usize, usize, usize),
    IndexSelect(Tensor<T>, IntTensor, usize),
    IndexAdd(Tensor<T>, IntTensor, Tensor<T>, usize),
    ScatterAdd(Tensor<T>, IntTensor, Tensor<T>, usize),
    Gather(Tensor<T>, IntTensor, usize),
    Reshape(Tensor<T>),
    Transpose(Tensor<T>, usize, usize),
    Permute(Tensor<T>, Vec<usize>),
    Cat(Vec<Tensor<T>>, usize),
    IfElse(Tensor<bool>, Option<Tensor<T>>, Option<Tensor<T>>),
    Copy(Tensor<T>),
    
    RmsNorm(Tensor<T>, Tensor<T>, T),
    Softmax(Tensor<T>, usize),
}

#[derive(Clone, Copy, PartialEq, Eq)]
pub enum CmpOp {
    Eq,
    Ne,
    Le,
    Ge,
    Lt,
    Gt,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReduceOp {
    Sum,
    Min,
    Max,
    Mean,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BinaryOp {
    Add,
    Mul,
    Sub,
    Div,
    Maximum,
    Minimum,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UnaryOp<T: WithDType> {
    Exp,
    Ln,

    Sin,
    Cos,
    Tanh,

    Abs,
    Neg,
    Sqr,
    Sqrt,

    Recip,
    Gelu,
    GeluErf,
    Erf,
    Relu,
    LeakyRelu(T),
    Silu,
    Sigmoid,

    Floor,
    Ceil,
    Round,
    Sign,

    Affine { mul: T, add: T },
    Pow(T),
}