luma_tensor/ops/
mod.rs

1//! The autograd computation-graph node type [`Op`] and its op-kind enums.
2//!
3//! Note on scalars: in the old design op scalars were typed `T` (tied to the
4//! tensor's element type). Here a `Float` tensor's precision is a runtime
5//! [`DType`](crate::DType), so all op scalars are stored as `f64` — the widest
6//! float — and cast back to the tensor's actual precision when the op runs or
7//! its gradient is computed.
8
9mod arith;
10mod boolean;
11mod cast;
12pub mod construct;
13mod display;
14mod indexer;
15mod matmul;
16mod nn;
17mod numeric;
18mod reduce;
19mod shape;
20mod transfer;
21
22pub use construct::{DEFAULT_FLOAT, DEFAULT_INT};
23pub use indexer::IndexingDTypeKind;
24pub use indexer::{IndexOp, Indexer, Slice};
25pub use numeric::NumericDTypeKind;
26pub use reduce::ReduceDTypeKind;
27pub use shape::ShapeDTypeKind;
28pub use transfer::TransferDTypeKind;
29
30use crate::{Bool, Device, Float, Int, Tensor};
31
32/// A node in the (implicit) computation graph: the operation that produced a
33/// `Float` tensor, holding `Arc` references to its inputs. Only `Float` tensors
34/// record ops — `Int`/`Bool` tensors are never differentiated.
35///
36/// Inputs are typed by kind: differentiable inputs are `Tensor<D, Float>`,
37/// index inputs are `Tensor<D, Int>`, and masks are `Tensor<D, Bool>`.
38pub enum Op<D: Device> {
39    Binary(Tensor<D, Float>, Tensor<D, Float>, BinaryOp),
40    BinaryScalarRhs(Tensor<D, Float>, f64, BinaryOp),
41    BinaryScalarLhs(f64, Tensor<D, Float>, BinaryOp),
42    FloatUnary(Tensor<D, Float>, FloatUnaryOp),
43    Unary(Tensor<D, Float>, UnaryOp<f64>),
44    Reduce(Tensor<D, Float>, ReduceOp, Vec<usize>),
45    Matmul(Tensor<D, Float>, Tensor<D, Float>),
46    Broadcast(Tensor<D, Float>),
47    Narrow(Tensor<D, Float>, usize, usize, usize),
48    Slice(Tensor<D, Float>, usize, usize, usize, usize),
49    IndexSelect(Tensor<D, Float>, Tensor<D, Int>, usize),
50    IndexAdd(Tensor<D, Float>, Tensor<D, Int>, Tensor<D, Float>, usize),
51    ScatterAdd(Tensor<D, Float>, Tensor<D, Int>, Tensor<D, Float>, usize),
52    Gather(Tensor<D, Float>, Tensor<D, Int>, usize),
53    Reshape(Tensor<D, Float>),
54    Transpose(Tensor<D, Float>, usize, usize),
55    Permute(Tensor<D, Float>, Vec<usize>),
56    Cat(Vec<Tensor<D, Float>>, usize),
57    Pick(Tensor<D, Bool>, Option<Tensor<D, Float>>, Option<Tensor<D, Float>>),
58    Copy(Tensor<D, Float>),
59    RmsNorm(Tensor<D, Float>, Tensor<D, Float>, f64),
60    Softmax(Tensor<D, Float>, usize),
61    /// Precision cast within the float kind (e.g. f32 -> f64). Records the input
62    /// so the gradient can be cast back — the key capability the old design lacked.
63    Cast(Tensor<D, Float>),
64}
65
66/// A *view* operation: it shares the source tensor's storage under a new
67/// [`Layout`] without touching device kernels.
68///
69/// Views bypass the `Device` kernel seam entirely (`Tensor::share_storage`), so
70/// tracing devices learn about them through [`Device::on_view`](crate::Device::on_view)
71/// rather than through `FloatOps`/`IntOps`/`BoolOps`.
72#[derive(Debug, Clone, PartialEq, Eq)]
73pub enum ViewOp {
74    Reshape,
75    Transpose(usize, usize),
76    Permute(Vec<usize>),
77    Narrow(usize, usize, usize),
78    Slice(usize, usize, usize, usize),
79    Broadcast,
80    Squeeze,
81    Unsqueeze,
82}
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub enum BinaryOp {
86    Add,
87    Mul,
88    Sub,
89    Div,
90    Maximum,
91    Minimum,
92}
93
94#[derive(Debug, Clone, Copy, PartialEq, Eq)]
95pub enum ReduceOp {
96    Sum,
97    Min,
98    Max,
99    Mean,
100    Prod,
101}
102
103#[derive(Debug, Clone, Copy, PartialEq)]
104pub enum FloatUnaryOp {
105    Exp,
106    Ln,
107    Sin,
108    Cos,
109    Tanh,
110    Sqr,
111    Sqrt,
112    Recip,
113    Gelu,
114    GeluErf,
115    Erf,
116    Relu,
117    LeakyRelu(f64),
118    Silu,
119    Sigmoid,
120    Floor,
121    Ceil,
122    Round,
123}
124
125#[derive(Debug, Clone, Copy, PartialEq)]
126pub enum UnaryOp<S> {
127    Neg,
128    Abs,
129    Sign,
130    Affine(S, S),
131    Pow(S),
132    Clamp(Option<S>, Option<S>),
133}
134
135#[derive(Debug, Clone, Copy, PartialEq, Eq)]
136pub enum CmpOp {
137    Eq,
138    Ne,
139    Le,
140    Ge,
141    Lt,
142    Gt,
143}
luma_tensor/ops/mod.rs

luma_tensor/ops/
mod.rs
