lumen_core/ops/
mod.rs

1mod construct;
2mod indexer;
3mod iter;
4pub mod display;
5mod shape;
6mod arith;
7mod matmul;
8mod reduce;
9mod broadcast;
10mod convert;
11mod boolean;
12mod fused;
13mod nn;
14mod variable;
15
16pub use construct::ToTensor;
17pub use indexer::{Slice, IndexOp};
18pub use iter::*;
19pub use indexer::*;
20pub use variable::*;
21
22use crate::{FloatDType, IntTensor, Tensor, WithDType};
23
24#[derive(Clone)]
25pub enum Op<T: FloatDType> {
26    Binary(Tensor<T>, Tensor<T>, BinaryOp),
27    BinaryScalarRhs(Tensor<T>, T, BinaryOp),
28    BinaryScalarLhs(T, Tensor<T>, BinaryOp),
29    Unary(Tensor<T>, UnaryOp<T>),
30    Reduce(Tensor<T>, ReduceOp, Vec<usize>),
31    Matmul(Tensor<T>, Tensor<T>),
32    Broadcast(Tensor<T>),
33    Narrow(Tensor<T>, usize, usize, usize),
34    Slice(Tensor<T>, usize, usize, usize, usize),
35    IndexSelect(Tensor<T>, IntTensor, usize),
36    IndexAdd(Tensor<T>, IntTensor, Tensor<T>, usize),
37    ScatterAdd(Tensor<T>, IntTensor, Tensor<T>, usize),
38    Gather(Tensor<T>, IntTensor, usize),
39    Reshape(Tensor<T>),
40    Transpose(Tensor<T>, usize, usize),
41    Permute(Tensor<T>, Vec<usize>),
42    Cat(Vec<Tensor<T>>, usize),
43    IfElse(Tensor<bool>, Option<Tensor<T>>, Option<Tensor<T>>),
44    Copy(Tensor<T>),
45    
46    RmsNorm(Tensor<T>, Tensor<T>, T),
47    Softmax(Tensor<T>, usize),
48}
49
50#[derive(Clone, Copy, PartialEq, Eq)]
51pub enum CmpOp {
52    Eq,
53    Ne,
54    Le,
55    Ge,
56    Lt,
57    Gt,
58}
59
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61pub enum ReduceOp {
62    Sum,
63    Min,
64    Max,
65    Mean,
66}
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq)]
69pub enum BinaryOp {
70    Add,
71    Mul,
72    Sub,
73    Div,
74    Maximum,
75    Minimum,
76}
77
78#[derive(Debug, Clone, Copy, PartialEq, Eq)]
79pub enum UnaryOp<T: WithDType> {
80    Exp,
81    Ln,
82
83    Sin,
84    Cos,
85    Tanh,
86
87    Abs,
88    Neg,
89    Sqr,
90    Sqrt,
91
92    Recip,
93    Gelu,
94    GeluErf,
95    Erf,
96    Relu,
97    LeakyRelu(T),
98    Silu,
99    Sigmoid,
100
101    Floor,
102    Ceil,
103    Round,
104    Sign,
105
106    Affine { mul: T, add: T },
107    Pow(T),
108}
lumen_core/ops/mod.rs

lumen_core/ops/
mod.rs
