luma_tensor/grad/
meta.rs

1use crate::{
2    BinaryOp, Bool, DTypeKind, Device, Float, FloatUnaryOp, Int, Op, ReduceOp, Tensor, TensorId, TensorImpl, UnaryOp, is_grad_enabled,
3};
4use std::sync::{Arc, RwLock};
5
6/// Trait for metadata that knows how to construct itself when a tensor operation is performed.
7pub trait TensorMeta<D: Device, K: DTypeKind<D> + Sized>: Default + Send + Sync {
8    // ---- Binary operations ----
9    fn on_binary(lhs: &Tensor<D, K>, rhs: &Tensor<D, K>, op: BinaryOp) -> Self;
10    fn on_binary_scalar(lhs: &Tensor<D, K>, rhs: K::Scalar, op: BinaryOp) -> Self;
11
12    // ---- Unary operations ----
13    fn on_unary(t: &Tensor<D, K>, op: UnaryOp<K::Scalar>) -> Self;
14    fn on_float_unary(t: &Tensor<D, K>, op: FloatUnaryOp) -> Self;
15
16    // ---- Reductions ----
17    fn on_reduce(t: &Tensor<D, K>, dims: &[usize], op: ReduceOp) -> Self;
18
19    // ---- Matrix operations ----
20    fn on_matmul(lhs: &Tensor<D, K>, rhs: &Tensor<D, K>) -> Self;
21
22    // ---- Shape operations ----
23    fn on_broadcast(t: &Tensor<D, K>) -> Self;
24    fn on_narrow(t: &Tensor<D, K>, dim: usize, start: usize, len: usize) -> Self;
25    fn on_slice(t: &Tensor<D, K>, dim: usize, start: usize, end: usize, step: usize) -> Self;
26    fn on_reshape(t: &Tensor<D, K>) -> Self;
27    fn on_transpose(t: &Tensor<D, K>, dim1: usize, dim2: usize) -> Self;
28    fn on_permute(t: &Tensor<D, K>, dims: Vec<usize>) -> Self;
29    fn on_cat<A: AsRef<Tensor<D, K>>>(args: &[A], dim: usize) -> Self;
30    fn on_copy(t: &Tensor<D, K>) -> Self;
31
32    // ---- Type conversions ----
33    fn on_cast(t: &Tensor<D, K>) -> Self;
34
35    // ---- Indexing operations ----
36    fn on_index_select(t: &Tensor<D, K>, idx: &Tensor<D, Int>, dim: usize) -> Self;
37    fn on_gather(src: &Tensor<D, K>, idx: &Tensor<D, Int>, dim: usize) -> Self;
38    fn on_index_add(init: &Tensor<D, K>, idx: &Tensor<D, Int>, src: &Tensor<D, K>, dim: usize) -> Self;
39    fn on_scatter_add(init: &Tensor<D, K>, idx: &Tensor<D, Int>, src: &Tensor<D, K>, dim: usize) -> Self;
40
41    // ---- Conditional operations ----
42    fn on_pick(mask: &Tensor<D, Bool>, tv: Option<&Tensor<D, K>>, fv: Option<&Tensor<D, K>>) -> Self;
43
44    // ---- NN operations (Float-specific, but included for completeness) ----
45    fn on_rms_norm(input: &Tensor<D, K>, weight: &Tensor<D, K>, eps: f64) -> Self;
46    fn on_softmax(input: &Tensor<D, K>, dim: usize) -> Self;
47}
48
49pub struct FloatMeta<D: Device> {
50    pub op: Option<Op<D>>,
51    pub requires_grad: RwLock<bool>,
52}
53
54impl<D: Device> FloatMeta<D> {
55    /// A leaf variable that accumulates gradients.
56    pub fn var() -> Self {
57        Self { op: None, requires_grad: RwLock::new(true) }
58    }
59
60    /// A constant value that does not track gradients.
61    pub fn val() -> Self {
62        Self { op: None, requires_grad: RwLock::new(false) }
63    }
64
65    /// A non-leaf node produced by `op`.
66    pub fn from_op(op: Op<D>) -> Self {
67        Self { op: Some(op), requires_grad: RwLock::new(true) }
68    }
69
70    pub fn op(&self) -> Option<&Op<D>> {
71        self.op.as_ref()
72    }
73
74    pub fn requires_grad(&self) -> bool {
75        *self.requires_grad.read().unwrap()
76    }
77
78    pub fn set_requires_grad(&self, mode: bool) {
79        *self.requires_grad.write().unwrap() = mode;
80    }
81
82    /// A tensor is a leaf when it requires grad but was not produced by an op.
83    pub fn is_leaf(&self) -> bool {
84        self.requires_grad() && self.op.is_none()
85    }
86}
87
88impl<D: Device> Default for FloatMeta<D> {
89    fn default() -> Self {
90        Self::val()
91    }
92}
93
94/// Records an op into a `FloatMeta` iff grad is globally enabled and `record` is
95/// true (i.e. some input requires grad). Otherwise produces a plain value meta.
96impl<D: Device> FloatMeta<D> {
97    fn record(record: bool, op: impl FnOnce() -> Op<D>) -> Self {
98        if is_grad_enabled() && record { Self::from_op(op()) } else { Self::val() }
99    }
100
101    pub fn on_binary(lhs: &Tensor<D, Float>, rhs: &Tensor<D, Float>, op: BinaryOp) -> Self {
102        Self::record(lhs.requires_grad() || rhs.requires_grad(), || Op::Binary(lhs.clone(), rhs.clone(), op))
103    }
104
105    pub fn on_binary_scalar(lhs: &Tensor<D, Float>, rhs: f64, op: BinaryOp) -> Self {
106        Self::record(lhs.requires_grad(), || Op::BinaryScalarRhs(lhs.clone(), rhs, op))
107    }
108
109    pub fn on_unary(t: &Tensor<D, Float>, op: UnaryOp<f64>) -> Self {
110        Self::record(t.requires_grad(), || Op::Unary(t.clone(), op))
111    }
112
113    pub fn on_float_unary(t: &Tensor<D, Float>, op: FloatUnaryOp) -> Self {
114        Self::record(t.requires_grad(), || Op::FloatUnary(t.clone(), op))
115    }
116
117    pub fn on_broadcast(t: &Tensor<D, Float>) -> Self {
118        Self::record(t.requires_grad(), || Op::Broadcast(t.clone()))
119    }
120
121    pub fn on_reduce(t: &Tensor<D, Float>, dims: &[usize], op: ReduceOp) -> Self {
122        Self::record(t.requires_grad(), || Op::Reduce(t.clone(), op, dims.to_vec()))
123    }
124
125    pub fn on_matmul(lhs: &Tensor<D, Float>, rhs: &Tensor<D, Float>) -> Self {
126        Self::record(lhs.requires_grad() || rhs.requires_grad(), || Op::Matmul(lhs.clone(), rhs.clone()))
127    }
128
129    pub fn on_narrow(t: &Tensor<D, Float>, dim: usize, start: usize, len: usize) -> Self {
130        Self::record(t.requires_grad(), || Op::Narrow(t.clone(), dim, start, len))
131    }
132
133    pub fn on_slice(t: &Tensor<D, Float>, dim: usize, start: usize, end: usize, step: usize) -> Self {
134        Self::record(t.requires_grad(), || Op::Slice(t.clone(), dim, start, end, step))
135    }
136
137    pub fn on_reshape(t: &Tensor<D, Float>) -> Self {
138        Self::record(t.requires_grad(), || Op::Reshape(t.clone()))
139    }
140
141    pub fn on_transpose(t: &Tensor<D, Float>, dim1: usize, dim2: usize) -> Self {
142        Self::record(t.requires_grad(), || Op::Transpose(t.clone(), dim1, dim2))
143    }
144
145    pub fn on_permute(t: &Tensor<D, Float>, dims: Vec<usize>) -> Self {
146        Self::record(t.requires_grad(), || Op::Permute(t.clone(), dims))
147    }
148
149    pub fn on_cat<A: AsRef<Tensor<D, Float>>>(args: &[A], dim: usize) -> Self {
150        let record = args.iter().any(|t| t.as_ref().requires_grad());
151        Self::record(record, || {
152            let vec = args.iter().map(|a| a.as_ref().clone()).collect();
153            Op::Cat(vec, dim)
154        })
155    }
156
157    pub fn on_copy(t: &Tensor<D, Float>) -> Self {
158        Self::record(t.requires_grad(), || Op::Copy(t.clone()))
159    }
160
161    pub fn on_cast(t: &Tensor<D, Float>) -> Self {
162        Self::record(t.requires_grad(), || Op::Cast(t.clone()))
163    }
164
165    pub fn on_pick(mask: &Tensor<D, Bool>, tv: Option<&Tensor<D, Float>>, fv: Option<&Tensor<D, Float>>) -> Self {
166        let record = tv.map(|t| t.requires_grad()).unwrap_or(false) || fv.map(|f| f.requires_grad()).unwrap_or(false);
167        Self::record(record, || Op::Pick(mask.clone(), tv.cloned(), fv.cloned()))
168    }
169
170    pub fn on_index_select(t: &Tensor<D, Float>, idx: &Tensor<D, Int>, dim: usize) -> Self {
171        Self::record(t.requires_grad(), || Op::IndexSelect(t.clone(), idx.clone(), dim))
172    }
173
174    pub fn on_index_add(init: &Tensor<D, Float>, idx: &Tensor<D, Int>, src: &Tensor<D, Float>, dim: usize) -> Self {
175        Self::record(init.requires_grad() || src.requires_grad(), || Op::IndexAdd(init.clone(), idx.clone(), src.clone(), dim))
176    }
177
178    pub fn on_scatter_add(init: &Tensor<D, Float>, idx: &Tensor<D, Int>, src: &Tensor<D, Float>, dim: usize) -> Self {
179        Self::record(init.requires_grad() || src.requires_grad(), || Op::ScatterAdd(init.clone(), idx.clone(), src.clone(), dim))
180    }
181
182    pub fn on_gather(src: &Tensor<D, Float>, idx: &Tensor<D, Int>, dim: usize) -> Self {
183        Self::record(src.requires_grad(), || Op::Gather(src.clone(), idx.clone(), dim))
184    }
185
186    pub fn on_rms_norm(input: &Tensor<D, Float>, weight: &Tensor<D, Float>, eps: f64) -> Self {
187        Self::record(input.requires_grad() || weight.requires_grad(), || Op::RmsNorm(input.clone(), weight.clone(), eps))
188    }
189
190    pub fn on_softmax(input: &Tensor<D, Float>, dim: usize) -> Self {
191        Self::record(input.requires_grad(), || Op::Softmax(input.clone(), dim))
192    }
193}
194
195/// Convenience accessors on a `Float` tensor for its autograd state.
196impl<D: Device> Tensor<D, Float> {
197    pub fn detach(&self) -> Self {
198        if !self.requires_grad() {
199            self.clone()
200        } else {
201            Self(Arc::new(TensorImpl {
202                id: TensorId::new(),
203                storage: self.0.storage.clone(),
204                layout: self.layout().clone(),
205                dtype: self.dtype(),
206                device: self.device().clone(),
207                meta: FloatMeta::val(),
208            }))
209        }
210    }
211
212    pub fn requires_grad(&self) -> bool {
213        self.0.meta.requires_grad()
214    }
215
216    pub fn set_requires_grad(&self, mode: bool) {
217        self.0.meta.set_requires_grad(mode)
218    }
219
220    pub fn is_leaf(&self) -> bool {
221        self.0.meta.is_leaf()
222    }
223
224    pub fn op(&self) -> Option<&Op<D>> {
225        self.0.meta.op()
226    }
227}
228
229// ============================================================================
230// TensorMeta trait implementation for FloatMeta
231// ============================================================================
232
233impl<D: Device> TensorMeta<D, Float> for FloatMeta<D> {
234    fn on_binary(lhs: &Tensor<D, Float>, rhs: &Tensor<D, Float>, op: BinaryOp) -> Self {
235        FloatMeta::on_binary(lhs, rhs, op)
236    }
237
238    fn on_binary_scalar(lhs: &Tensor<D, Float>, rhs: f64, op: BinaryOp) -> Self {
239        FloatMeta::on_binary_scalar(lhs, rhs, op)
240    }
241
242    fn on_unary(t: &Tensor<D, Float>, op: UnaryOp<f64>) -> Self {
243        FloatMeta::on_unary(t, op)
244    }
245
246    fn on_float_unary(t: &Tensor<D, Float>, op: FloatUnaryOp) -> Self {
247        FloatMeta::on_float_unary(t, op)
248    }
249
250    fn on_reduce(t: &Tensor<D, Float>, dims: &[usize], op: ReduceOp) -> Self {
251        FloatMeta::on_reduce(t, dims, op)
252    }
253
254    fn on_matmul(lhs: &Tensor<D, Float>, rhs: &Tensor<D, Float>) -> Self {
255        FloatMeta::on_matmul(lhs, rhs)
256    }
257
258    fn on_broadcast(t: &Tensor<D, Float>) -> Self {
259        FloatMeta::on_broadcast(t)
260    }
261
262    fn on_narrow(t: &Tensor<D, Float>, dim: usize, start: usize, len: usize) -> Self {
263        FloatMeta::on_narrow(t, dim, start, len)
264    }
265
266    fn on_slice(t: &Tensor<D, Float>, dim: usize, start: usize, end: usize, step: usize) -> Self {
267        FloatMeta::on_slice(t, dim, start, end, step)
268    }
269
270    fn on_reshape(t: &Tensor<D, Float>) -> Self {
271        FloatMeta::on_reshape(t)
272    }
273
274    fn on_transpose(t: &Tensor<D, Float>, dim1: usize, dim2: usize) -> Self {
275        FloatMeta::on_transpose(t, dim1, dim2)
276    }
277
278    fn on_permute(t: &Tensor<D, Float>, dims: Vec<usize>) -> Self {
279        FloatMeta::on_permute(t, dims)
280    }
281
282    fn on_cat<A: AsRef<Tensor<D, Float>>>(args: &[A], dim: usize) -> Self {
283        FloatMeta::on_cat(args, dim)
284    }
285
286    fn on_copy(t: &Tensor<D, Float>) -> Self {
287        FloatMeta::on_copy(t)
288    }
289
290    fn on_cast(t: &Tensor<D, Float>) -> Self {
291        FloatMeta::on_cast(t)
292    }
293
294    fn on_index_select(t: &Tensor<D, Float>, idx: &Tensor<D, Int>, dim: usize) -> Self {
295        FloatMeta::on_index_select(t, idx, dim)
296    }
297
298    fn on_gather(src: &Tensor<D, Float>, idx: &Tensor<D, Int>, dim: usize) -> Self {
299        FloatMeta::on_gather(src, idx, dim)
300    }
301
302    fn on_index_add(init: &Tensor<D, Float>, idx: &Tensor<D, Int>, src: &Tensor<D, Float>, dim: usize) -> Self {
303        FloatMeta::on_index_add(init, idx, src, dim)
304    }
305
306    fn on_scatter_add(init: &Tensor<D, Float>, idx: &Tensor<D, Int>, src: &Tensor<D, Float>, dim: usize) -> Self {
307        FloatMeta::on_scatter_add(init, idx, src, dim)
308    }
309
310    fn on_pick(mask: &Tensor<D, Bool>, tv: Option<&Tensor<D, Float>>, fv: Option<&Tensor<D, Float>>) -> Self {
311        FloatMeta::on_pick(mask, tv, fv)
312    }
313
314    fn on_rms_norm(input: &Tensor<D, Float>, weight: &Tensor<D, Float>, eps: f64) -> Self {
315        FloatMeta::on_rms_norm(input, weight, eps)
316    }
317
318    fn on_softmax(input: &Tensor<D, Float>, dim: usize) -> Self {
319        FloatMeta::on_softmax(input, dim)
320    }
321}
322
323// ============================================================================
324// TensorMeta trait implementation for () (Int and Bool)
325// ============================================================================
326
327impl<D: Device, K: crate::DTypeKind<D>> TensorMeta<D, K> for () {
328    fn on_binary(_: &Tensor<D, K>, _: &Tensor<D, K>, _: BinaryOp) -> Self {}
329    fn on_binary_scalar(_: &Tensor<D, K>, _: K::Scalar, _: BinaryOp) -> Self {}
330    fn on_unary(_: &Tensor<D, K>, _: UnaryOp<K::Scalar>) -> Self {}
331    fn on_float_unary(_: &Tensor<D, K>, _: FloatUnaryOp) -> Self {}
332    fn on_reduce(_: &Tensor<D, K>, _: &[usize], _: ReduceOp) -> Self {}
333    fn on_matmul(_: &Tensor<D, K>, _: &Tensor<D, K>) -> Self {}
334    fn on_broadcast(_: &Tensor<D, K>) -> Self {}
335    fn on_narrow(_: &Tensor<D, K>, _: usize, _: usize, _: usize) -> Self {}
336    fn on_slice(_: &Tensor<D, K>, _: usize, _: usize, _: usize, _: usize) -> Self {}
337    fn on_reshape(_: &Tensor<D, K>) -> Self {}
338    fn on_transpose(_: &Tensor<D, K>, _: usize, _: usize) -> Self {}
339    fn on_permute(_: &Tensor<D, K>, _: Vec<usize>) -> Self {}
340    fn on_cat<A: AsRef<Tensor<D, K>>>(_: &[A], _: usize) -> Self {}
341    fn on_copy(_: &Tensor<D, K>) -> Self {}
342    fn on_cast(_: &Tensor<D, K>) -> Self {}
343    fn on_index_select(_: &Tensor<D, K>, _: &Tensor<D, Int>, _: usize) -> Self {}
344    fn on_gather(_: &Tensor<D, K>, _: &Tensor<D, Int>, _: usize) -> Self {}
345    fn on_index_add(_: &Tensor<D, K>, _: &Tensor<D, Int>, _: &Tensor<D, K>, _: usize) -> Self {}
346    fn on_scatter_add(_: &Tensor<D, K>, _: &Tensor<D, Int>, _: &Tensor<D, K>, _: usize) -> Self {}
347    fn on_pick(_: &Tensor<D, Bool>, _: Option<&Tensor<D, K>>, _: Option<&Tensor<D, K>>) -> Self {}
348    fn on_rms_norm(_: &Tensor<D, K>, _: &Tensor<D, K>, _: f64) -> Self {}
349    fn on_softmax(_: &Tensor<D, K>, _: usize) -> Self {}
350}
luma_tensor/grad/meta.rs

luma_tensor/grad/
meta.rs
