lumen_core/grad/backprop.rs
1use std::collections::HashMap;
2use crate::{FloatDType, Tensor, TensorId,};
3use crate::{BinaryOp, ReduceOp, UnaryOp, Op};
4use super::GradStore;
5
6impl<T: FloatDType> Tensor<T> {
7
8 pub fn backward(&self) -> crate::Result<GradStore<T>> {
9 let _guard = crate::NoGradGuard::new();
10
11 let sorted_nodes = self.sorted_nodes();
12 let mut grads = GradStore::new();
13 grads.insert(self, self.ones_like()?);
14
15 for node in sorted_nodes.iter() {
16 match node.op() {
17 None => {
18 assert!(node.is_leaf());
19 continue
20 }
21 Some(op) => {
22 let grad = grads
23 .remove(node)
24 .expect("grad not populated");
25
26 match op {
27 //=========================================================================================//
28 // Binary
29 //=========================================================================================//
30 Op::Binary(lhs, rhs, BinaryOp::Add) => {
31 let lhs_sum_grad = grads.or_insert(lhs)?;
32 lhs_sum_grad.impl_add_(&grad)?;
33 let rhs_sum_grad = grads.or_insert(rhs)?;
34 rhs_sum_grad.impl_add_(&grad)?;
35 }
36 Op::Binary(lhs, rhs, BinaryOp::Sub) => {
37 let lhs_sum_grad = grads.or_insert(lhs)?;
38 lhs_sum_grad.impl_add_(&grad)?;
39 let rhs_sum_grad = grads.or_insert(rhs)?;
40 rhs_sum_grad.impl_sub_(&grad)?;
41 }
42 Op::Binary(lhs, rhs, BinaryOp::Mul) => {
43 let lhs_grad = grad.mul(rhs)?;
44 let lhs_sum_grad = grads.or_insert(lhs)?;
45 lhs_sum_grad.impl_add_(&lhs_grad)?;
46
47 let rhs_grad = grad.mul(lhs)?;
48 let rhs_sum_grad = grads.or_insert(rhs)?;
49 rhs_sum_grad.impl_add_(&rhs_grad)?;
50 }
51 Op::Binary(lhs, rhs, BinaryOp::Div) => {
52 let lhs_grad = grad.div(rhs)?;
53 let lhs_sum_grad = grads.or_insert(lhs)?;
54 lhs_sum_grad.impl_add_(&lhs_grad)?;
55
56 let rhs_grad = grad.mul(lhs)?.div(&rhs.sqr()?)?;
57 let rhs_sum_grad = grads.or_insert(rhs)?;
58 rhs_sum_grad.impl_sub_(&rhs_grad)?;
59 }
60 Op::Binary(lhs, rhs, BinaryOp::Minimum)
61 | Op::Binary(lhs, rhs, BinaryOp::Maximum) => {
62 let mask_lhs = (*node).eq(lhs)?.cast()?;
63 let mask_rhs = (*node).eq(rhs)?.cast()?;
64
65 // If both masks are 1 one the same point, we want to scale the
66 // gradient by 0.5 rather than 1.
67 let lhs_grad = mask_lhs.mul(&grad)?.div(&(&mask_rhs + T::one()))?;
68 let lhs_sum_grad = grads.or_insert(lhs)?;
69 lhs_sum_grad.impl_add_(&lhs_grad)?;
70
71 let rhs_grad = mask_rhs.mul(&grad)?.div(&(&mask_lhs + T::one()))?;
72 let rhs_sum_grad = grads.or_insert(rhs)?;
73 rhs_sum_grad.impl_add_(&rhs_grad)?;
74 }
75
76 //=========================================================================================//
77 // BinaryScalarRhs
78 //=========================================================================================//
79 Op::BinaryScalarRhs(lhs, _, BinaryOp::Add) => {
80 // y = x + c => dy/dx = 1
81 let lhs_sum_grad = grads.or_insert(lhs)?;
82 lhs_sum_grad.impl_add_(&grad)?;
83 }
84 Op::BinaryScalarRhs(lhs, _, BinaryOp::Sub) => {
85 // y = x - c => dy/dx = 1
86 let lhs_sum_grad = grads.or_insert(lhs)?;
87 lhs_sum_grad.impl_add_(&grad)?;
88 }
89 Op::BinaryScalarRhs(lhs, rhs, BinaryOp::Mul) => {
90 // y = x * c => dy/dx = c
91 let lhs_grad = grad.mul_scalar(*rhs)?;
92 let lhs_sum_grad = grads.or_insert(lhs)?;
93 lhs_sum_grad.impl_add_(&lhs_grad)?;
94 }
95 Op::BinaryScalarRhs(lhs, rhs, BinaryOp::Div) => {
96 // y = x / c => dy/dx = 1/c
97 let lhs_grad = grad.div_scalar(*rhs)?;
98 let lhs_sum_grad = grads.or_insert(lhs)?;
99 lhs_sum_grad.impl_add_(&lhs_grad)?;
100 }
101 Op::BinaryScalarRhs(lhs, rhs, BinaryOp::Maximum) |
102 Op::BinaryScalarRhs(lhs, rhs, BinaryOp::Minimum) => {
103 let mask_lhs = (*node).eq(lhs)?.cast()?;
104 let mask_rhs = (*node).eq(*rhs)?.cast()?;
105 let lhs_grad = mask_lhs.mul(&grad)?.div(&(&mask_rhs + T::one()))?;
106 let lhs_sum_grad = grads.or_insert(lhs)?;
107 lhs_sum_grad.impl_add_(&lhs_grad)?;
108 }
109
110 //=========================================================================================//
111 // BinaryScalarLhs
112 //=========================================================================================//
113 Op::BinaryScalarLhs(_, rhs, BinaryOp::Add) => {
114 // y = c + x => dy/dx = 1
115 let rhs_sum_grad = grads.or_insert(rhs)?;
116 rhs_sum_grad.impl_add_(&grad)?;
117 }
118 Op::BinaryScalarLhs(_, rhs, BinaryOp::Sub) => {
119 // y = c - x => dy/dx = -1
120 let rhs_sum_grad = grads.or_insert(rhs)?;
121 rhs_sum_grad.impl_sub_(&grad)?;
122 }
123 Op::BinaryScalarLhs(lhs, rhs, BinaryOp::Mul) => {
124 // y = c * x => dy/dx = c
125 let rhs_grad = grad.mul_scalar(*lhs)?;
126 let rhs_sum_grad = grads.or_insert(rhs)?;
127 rhs_sum_grad.impl_add_(&rhs_grad)?;
128 }
129 Op::BinaryScalarLhs(lhs, rhs, BinaryOp::Div) => {
130 // y = c / x = c * x^(-1)
131 // dy/dx = -c * x^(-2) = -c / (x^2)
132 // grad_input = grad * (-c / x^2)
133 let numerator = grad.mul_scalar(-*lhs)?;
134 let denominator = rhs.mul(rhs)?;
135 let rhs_grad = numerator.div(&denominator)?;
136
137 let rhs_sum_grad = grads.or_insert(rhs)?;
138 rhs_sum_grad.impl_add_(&rhs_grad)?;
139 }
140 Op::BinaryScalarLhs(lhs, rhs, BinaryOp::Maximum) |
141 Op::BinaryScalarLhs(lhs, rhs, BinaryOp::Minimum) => {
142 let mask_lhs = (*node).eq(*lhs)?.cast()?;
143 let mask_rhs = (*node).eq(rhs)?.cast()?;
144 let rhs_grad = mask_rhs.mul(&grad)?.div(&(&mask_lhs + T::one()))?;
145 let rhs_sum_grad = grads.or_insert(rhs)?;
146 rhs_sum_grad.impl_add_(&rhs_grad)?;
147 }
148
149 //=========================================================================================//
150 // Unary
151 //=========================================================================================//
152 Op::Unary(_, UnaryOp::Ceil) => Err(crate::Error::BackwardNotSupported("ceil"))?,
153 Op::Unary(_, UnaryOp::Floor) => Err(crate::Error::BackwardNotSupported("floor"))?,
154 Op::Unary(_, UnaryOp::Round) => Err(crate::Error::BackwardNotSupported("round"))?,
155 Op::Unary(_, UnaryOp::Sign) => Err(crate::Error::BackwardNotSupported("sign"))?,
156 Op::Unary(arg, UnaryOp::Exp) => {
157 let arg_grad = grads.or_insert(arg)?;
158 arg_grad.impl_add_mul_(&grad, node)?;
159 }
160 Op::Unary(arg, UnaryOp::Ln) => {
161 let sum_grad = grads.or_insert(arg)?;
162 sum_grad.impl_add_div_(&grad, arg)?;
163 }
164 Op::Unary(arg, UnaryOp::Sin) => {
165 let sum_grad = grads.or_insert(arg)?;
166 sum_grad.fused_2arg_inplace_op(
167 &grad, &arg, T::add, |a, b| a * b.cos(), "fused in sin backward"
168 )?;
169 }
170 Op::Unary(arg, UnaryOp::Cos) => {
171 let sum_grad = grads.or_insert(arg)?;
172 sum_grad.fused_2arg_inplace_op(
173 &grad, &arg, T::sub, |a, b| a * b.sin(), "fused in Cos backward"
174 )?;
175 }
176 Op::Unary(arg, UnaryOp::Tanh) => {
177 let sum_grad = grads.or_insert(arg)?;
178 sum_grad.fused_2arg_inplace_op(
179 &grad, &node, T::sub, |a, b| a * b.sqr() - T::ONE, "fused in Tanh backward"
180 )?;
181 }
182 Op::Unary(arg, UnaryOp::Sqr) => {
183 let sum_grad = grads.or_insert(arg)?;
184 sum_grad.fused_2arg_inplace_op(
185 &arg, &grad,
186 |a, b| a + T::two() * b, T::mul, "fused in Sqr backward"
187 )?;
188 }
189 Op::Unary(arg, UnaryOp::Sqrt) => {
190 let sum_grad = grads.or_insert(arg)?;
191 sum_grad.fused_2arg_inplace_op(
192 &grad, &node,
193 |a, b| a + T::half() * b, T::div, "fused in Sqrt backward"
194 )?;
195 }
196 Op::Unary(arg, UnaryOp::Abs) => {
197 let sum_grad = grads.or_insert(arg)?;
198 sum_grad.fused_2arg_inplace_op(
199 &grad, &arg,
200 T::add, |a, b| a * ( if b >= T::ZERO {T::ONE} else {T::ONE.neg()} ) , "fused in Abs backward"
201 )?;
202 }
203 Op::Unary(arg, UnaryOp::Neg) => {
204 let sum_grad = grads.or_insert(arg)?;
205 sum_grad.impl_sub_(&grad)?;
206 }
207 Op::Unary(arg, UnaryOp::Recip) => {
208 let sum_grad = grads.or_insert(arg)?;
209 sum_grad.fused_2arg_inplace_op(
210 &grad, &arg,
211 T::sub, |a, b| a / b.sqr(), "fused in Recip backward"
212 )?;
213 }
214 Op::Unary(arg, UnaryOp::Gelu) => {
215 let sum_grad = grads.or_insert(arg)?;
216 let c1 = T::from_f64(0.0356774);
217 let c2 = T::from_f64(0.797885);
218 let c3 = T::from_f64(0.0535161);
219 let c4 = T::from_f64(0.398942);
220
221 sum_grad.fused_2arg_inplace_op(
222 &grad, &arg,
223 T::add,
224 move |g, x| {
225 let x3 = x * x.sqr();
226 let inner = x3 * c1 + x * c2;
227 let tanh_val = inner.tanh();
228 let dt = x3 * c3 + x * c4;
229 let gelu_grad = tanh_val * T::half() + dt * (T::one() - tanh_val.sqr()) + T::half();
230 g * gelu_grad
231 },
232 "fused in Gelu backward"
233 )?;
234 }
235 Op::Unary(arg, UnaryOp::Erf) => {
236 let sum_grad = grads.or_insert(arg)?;
237 let scale = T::two() / T::pi().sqrt();
238 sum_grad.fused_2arg_inplace_op(
239 &grad, &arg,
240 T::add,
241 move |g, x| g * x.sqr().neg().exp() * scale,
242 "fused in Erf backward"
243 )?;
244 }
245 Op::Unary(arg, UnaryOp::GeluErf) => {
246 let sum_grad = grads.or_insert(arg)?;
247 let c1 = T::from_f64(0.398942);
248 let sqrt_two = T::two().sqrt();
249
250 sum_grad.fused_2arg_inplace_op(
251 &grad, &arg,
252 T::add,
253 move |g, x| {
254 let neg_half_square = x.sqr().neg() / T::two();
255 let scaled_exp_arg = c1 * neg_half_square.exp() * x;
256 let arg_scaled_sqrt = x / sqrt_two;
257 let erf_scaled_sqrt = arg_scaled_sqrt.erf() / T::two();
258 let gelu_erf_grad = scaled_exp_arg + erf_scaled_sqrt + T::half();
259 g * gelu_erf_grad
260 },
261 "fused in GeluErf backward"
262 )?;
263 }
264 Op::Unary(arg, UnaryOp::Relu) => {
265 let sum_grad = grads.or_insert(arg)?;
266 sum_grad.fused_2arg_inplace_op(
267 &grad, &arg,
268 T::add,
269 |g, x| if x >= T::zero() { g } else { T::zero() },
270 "fused in Relu backward"
271 )?;
272 }
273 Op::Unary(arg, UnaryOp::Silu) => {
274 let sum_grad = grads.or_insert(arg)?;
275 sum_grad.fused_2arg_inplace_op(
276 &grad, &arg,
277 T::add,
278 |g, x| {
279 let sigmoid = (x.neg().exp() + T::one()).recip();
280 let silu_val = x * sigmoid;
281 let silu_grad = sigmoid * (T::one() - silu_val) + silu_val;
282 g * silu_grad
283 },
284 "fused in Silu backward"
285 )?;
286 }
287 Op::Unary(arg, UnaryOp::Sigmoid) => {
288 let sum_grad = grads.or_insert(arg)?;
289 sum_grad.fused_2arg_inplace_op(
290 &grad, &node,
291 T::add,
292 |g, y| g * y * (T::one() - y),
293 "fused in Sigmoid backward"
294 )?;
295 }
296 Op::Unary(arg, UnaryOp::LeakyRelu(negative_slope)) => {
297 let sum_grad = grads.or_insert(arg)?;
298 sum_grad.fused_2arg_inplace_op(
299 &grad, &arg,
300 T::add,
301 move |g, x| if x >= T::zero() { g } else { g * *negative_slope },
302 "fused in LeakyRelu backward"
303 )?;
304 }
305 Op::Unary(arg, UnaryOp::Pow(e)) => {
306 let sum_grad = grads.or_insert(arg)?;
307 sum_grad.fused_2arg_inplace_op(
308 &grad, &arg,
309 T::add,
310 move |g, x| g * x.powf(*e - T::one()) * *e,
311 "fused in Pow backward"
312 )?;
313 }
314 Op::Unary(arg, UnaryOp::Affine { mul, add: _ }) => {
315 // y = x * mul + add
316 // dy = mul
317 let arg_grad = grad.mul_scalar(*mul)?;
318 let lhs_sum_grad = grads.or_insert(arg)?;
319 lhs_sum_grad.impl_add_(&arg_grad)?;
320 }
321
322 //=========================================================================================//
323 // Matmul
324 //=========================================================================================//
325 Op::Matmul(lhs, rhs) => {
326 let lhs_sum_grad = grads.or_insert(lhs)?;
327 lhs_sum_grad.add_matmul_(&grad, &rhs.transpose_last()?)?;
328
329 let rhs_sum_grad = grads.or_insert(rhs)?;
330 rhs_sum_grad.add_matmul_(&lhs.transpose_last()?, &grad)?;
331 }
332
333 //=========================================================================================//
334 // Reduce
335 //=========================================================================================//
336 Op::Reduce(arg, ReduceOp::Sum, reduced_dims) => {
337 let grad = Self::broadcast_back(arg, &grad, reduced_dims)?;
338 let sum_grad = grads.or_insert(arg)?;
339 sum_grad.impl_add_(&grad)?;
340 }
341 Op::Reduce(arg, ReduceOp::Max, reduced_dims) => {
342 let node = Self::broadcast_back(arg, node, reduced_dims)?;
343 let grad = Self::broadcast_back(arg, &grad, reduced_dims)?;
344 let grad = node.eq(arg)?.cast()?.mul(&grad)?;
345 let sum_grad = grads.or_insert(arg)?;
346 sum_grad.impl_add_(&grad.broadcast_as(sum_grad.dims())?)?;
347 }
348 Op::Reduce(arg, ReduceOp::Min, reduced_dims) => {
349 let node = Self::broadcast_back(arg, node, reduced_dims)?;
350 let grad = Self::broadcast_back(arg, &grad, reduced_dims)?;
351 let grad = node.eq(arg)?.cast()?.mul(&grad)?;
352 let sum_grad = grads.or_insert(arg)?;
353 sum_grad.impl_add_(&grad.broadcast_as(sum_grad.dims())?)?;
354 }
355 Op::Reduce(arg, ReduceOp::Mean, reduced_dims) => {
356 let grad_output = Self::broadcast_back(arg, &grad, reduced_dims)?;
357 let n = arg.element_count() / node.element_count();
358 let scale = T::from_usize(n);
359
360 let sum_grad = grads.or_insert(arg)?;
361 sum_grad.fused_1arg_inplace_op(
362 &grad_output, T::add, |v| v / scale , "fused in Mean backward"
363 )?;
364 }
365
366 //=========================================================================================//
367 // Broadcast
368 //=========================================================================================//
369 Op::Broadcast(arg) => {
370 let arg_dims = arg.dims();
371 let node_dims = node.dims();
372 let left_dims = node_dims.len() - arg_dims.len();
373 let mut sum_dims: Vec<usize> = (0..left_dims).collect();
374 for (dim, (node_dim, arg_dim)) in node_dims[left_dims..]
375 .iter()
376 .zip(arg_dims.iter())
377 .enumerate()
378 {
379 if node_dim != arg_dim {
380 sum_dims.push(dim + left_dims)
381 }
382 }
383
384 let mut arg_grad = grad;
385 for &dim in sum_dims.iter() {
386 arg_grad = arg_grad.sum_keepdim(dim)?;
387 }
388
389 for _i in 0..left_dims {
390 arg_grad = arg_grad.squeeze(0)?
391 }
392 let sum_grad = grads.or_insert(arg)?;
393 sum_grad.impl_add_(&arg_grad.broadcast_as(sum_grad.dims())?)?;
394 }
395
396 //=========================================================================================//
397 // Narrow
398 //=========================================================================================//
399 &Op::Narrow(ref arg, dim, start_idx, len) => {
400 let arg_dims = arg.dims();
401 let left_pad = if start_idx == 0 {
402 None
403 } else {
404 let mut dims = arg_dims.to_vec();
405 dims[dim] = start_idx;
406 Some(Tensor::zeros(dims)?)
407 };
408 let right_pad = arg_dims[dim] - start_idx - len;
409 let right_pad = if right_pad == 0 {
410 None
411 } else {
412 let mut dims = arg_dims.to_vec();
413 dims[dim] = right_pad;
414 Some(Tensor::zeros(dims)?)
415 };
416 let arg_grad = match (left_pad, right_pad) {
417 (None, None) => grad,
418 (Some(l), None) => Tensor::cat(&[&l, &grad], dim)?,
419 (None, Some(r)) => Tensor::cat(&[&grad, &r], dim)?,
420 (Some(l), Some(r)) => Tensor::cat(&[&l, &grad, &r], dim)?,
421 };
422 let sum_grad = grads.or_insert(arg)?;
423 sum_grad.impl_add_(&arg_grad)?;
424 }
425
426 //=========================================================================================//
427 // Slice
428 //=========================================================================================//
429 &Op::Slice(ref arg, dim, start, _end, step) => {
430 let arg_dims = arg.dims();
431
432 let body_grad = if step == 1 {
433 grad
434 } else {
435 let grad_len = grad.dims()[dim];
436 let span_len = if grad_len > 0 { (grad_len - 1) * step + 1 } else { 0 };
437
438 let mut unsqueezed_shape = grad.dims().to_vec();
439 unsqueezed_shape.insert(dim + 1, 1);
440 let grad_unsqueezed = grad.reshape(&unsqueezed_shape)?;
441
442 let mut zeros_shape = unsqueezed_shape.clone();
443 zeros_shape[dim + 1] = step - 1;
444 let zeros_gap = Tensor::zeros(zeros_shape)?;
445
446 let dilated = Tensor::cat(&[&grad_unsqueezed, &zeros_gap], dim + 1)?;
447
448 let mut flattened_shape = grad.dims().to_vec();
449 flattened_shape[dim] = grad_len * step;
450 let flattened = dilated.reshape(flattened_shape)?;
451
452 flattened.narrow(dim, 0, span_len)?
453 };
454
455 let body_len = body_grad.dims()[dim];
456
457 let left_pad = if start == 0 {
458 None
459 } else {
460 let mut dims = arg_dims.to_vec();
461 dims[dim] = start;
462 Some(Tensor::zeros(dims)?)
463 };
464
465 let right_pad_len = arg_dims[dim] - start - body_len;
466 let right_pad = if right_pad_len == 0 {
467 None
468 } else {
469 let mut dims = arg_dims.to_vec();
470 dims[dim] = right_pad_len;
471 Some(Tensor::zeros(dims)?)
472 };
473
474 let arg_grad = match (left_pad, right_pad) {
475 (None, None) => body_grad,
476 (Some(l), None) => Tensor::cat(&[&l, &body_grad], dim)?,
477 (None, Some(r)) => Tensor::cat(&[&body_grad, &r], dim)?,
478 (Some(l), Some(r)) => Tensor::cat(&[&l, &body_grad, &r], dim)?,
479 };
480
481 let sum_grad = grads.or_insert(arg)?;
482 sum_grad.impl_add_(&arg_grad)?;
483 }
484
485 //=========================================================================================//
486 // Reshape
487 //=========================================================================================//
488 Op::Reshape(arg) => {
489 let arg_grad = grad.reshape(arg.dims())?;
490 let sum_grad = grads.or_insert(arg)?;
491 sum_grad.impl_add_(&arg_grad)?;
492 }
493
494 //=========================================================================================//
495 // Transpose
496 //=========================================================================================//
497 Op::Transpose(arg, dim1, dim2) => {
498 let arg_grad = grad.transpose(*dim1, *dim2)?;
499 let sum_grad = grads.or_insert(arg)?;
500 sum_grad.impl_add_(&arg_grad)?;
501 }
502
503 //=========================================================================================//
504 // Permute
505 //=========================================================================================//
506 Op::Permute(arg, dims) => {
507 let mut inv_dims = vec![0; dims.len()];
508 for (i, &dim_idx) in dims.iter().enumerate() {
509 inv_dims[dim_idx] = i
510 }
511 let arg_grad = grad.permute(inv_dims)?;
512 let sum_grad = grads.or_insert(arg)?;
513 sum_grad.impl_add_(&arg_grad)?;
514 }
515
516 //=========================================================================================//
517 // Cat
518 //=========================================================================================//
519 Op::Cat(args, dim) => {
520 let mut start_idx = 0;
521 for arg in args {
522 let len = arg.dims()[*dim];
523 let arg_grad = grad.narrow(*dim, start_idx, len)?;
524 let sum_grad = grads.or_insert(arg)?;
525 sum_grad.impl_add_(&arg_grad)?;
526 start_idx += len;
527 }
528 }
529
530 //=========================================================================================//
531 // Copy
532 //=========================================================================================//
533 Op::Copy(arg) => {
534 let sum_grad = grads.or_insert(arg)?;
535 sum_grad.impl_add_(&grad)?;
536 }
537
538 //=========================================================================================//
539 // IfElse
540 //=========================================================================================//
541 Op::IfElse(mask, tv, fv) => {
542 if let Some(tv) = tv {
543 let masked_grad = mask.if_else(&grad, T::zero())?;
544 let sum_grad = grads.or_insert(tv)?;
545 sum_grad.impl_add_(&masked_grad)?;
546 }
547
548 if let Some(fv) = fv {
549 let masked_grad = mask.if_else(T::zero(), &grad)?;
550 let sum_grad = grads.or_insert(fv)?;
551 sum_grad.impl_add_(&masked_grad)?;
552 }
553 }
554
555 //=========================================================================================//
556 // IndexSelect
557 //=========================================================================================//
558 Op::IndexSelect(arg, indexes, dim) => {
559 let sum_grad = grads.or_insert(arg)?;
560 *sum_grad = sum_grad.index_add(indexes.clone(), &grad, *dim)?;
561 }
562
563 //=========================================================================================//
564 // IndexAdd
565 //=========================================================================================//
566 Op::IndexAdd(init, indexes, src, dim) => {
567 let init_sum_grad = grads.or_insert(init)?;
568 *init_sum_grad = init_sum_grad.add(&grad)?;
569
570 let src_grad = grad.index_select(indexes.clone(), *dim)?;
571 let src_sum_grad = grads.or_insert(src)?;
572 *src_sum_grad = src_sum_grad.add(&src_grad)?;
573 }
574
575 //=========================================================================================//
576 // IndexAdd
577 //=========================================================================================//
578 #[allow(unused)]
579 Op::ScatterAdd(init, indexes, src, dim) => {
580 unimplemented!()
581 }
582
583 //=========================================================================================//
584 // Gather
585 //=========================================================================================//
586 Op::Gather(arg, indexes, dim) => {
587 let arg_grad = grads.or_insert(arg)?;
588 *arg_grad = arg_grad.scatter_add(indexes.clone(), &grad, *dim)?;
589 }
590
591 //=========================================================================================//
592 // Nn
593 //=========================================================================================//
594 Op::RmsNorm(input, weight, eps) => {
595 let input_grad = grads.or_insert(input)?.clone();
596 let weight_grad = grads.or_insert(weight)?;
597 Tensor::rms_norm_backward_fused(&grad, input, weight, &input_grad, &weight_grad, *eps)?;
598 }
599
600 Op::Softmax(input, dim) => {
601 let input_grad = grads.or_insert(input)?.clone();
602 Tensor::softmax_backward_fused(&grad, &node, &input_grad, *dim)?;
603 }
604 }
605 }
606 }
607 }
608
609 Ok(grads)
610 }
611
612 pub fn sorted_nodes(&self) -> Vec<&Tensor<T>> {
613 // The vec of sorted nodes is passed as an owned value rather than a mutable reference
614 // to get around some lifetime limitations.
615 fn walk<'a, T: FloatDType>(
616 node: &'a Tensor<T>,
617 nodes: Vec<&'a Tensor<T>>,
618 already_seen: &mut HashMap<TensorId, bool>,
619 ) -> (bool, Vec<&'a Tensor<T>>) {
620 if let Some(&tg) = already_seen.get(&node.id()) {
621 return (tg, nodes);
622 }
623 let mut track_grad = false;
624 let mut nodes = if node.is_leaf() {
625 track_grad = true;
626 nodes
627 } else if node.dtype().is_int() {
628 nodes
629 } else if let Some(op) = node.op() {
630 match op {
631 | Op::Binary(lhs, rhs, _)
632 | Op::Matmul(lhs, rhs)
633 | Op::IfElse(_, Some(lhs), Some(rhs))
634 | Op::IndexAdd(lhs, _, rhs, _)
635 | Op::ScatterAdd(lhs, _, rhs, _)
636 => {
637 let (tg, nodes) = walk(lhs, nodes, already_seen);
638 track_grad |= tg;
639 let (tg, nodes) = walk(rhs, nodes, already_seen);
640 track_grad |= tg;
641 nodes
642 }
643
644 | Op::Unary(_node, UnaryOp::Ceil)
645 | Op::Unary(_node, UnaryOp::Floor)
646 | Op::Unary(_node, UnaryOp::Round)
647 | Op::Unary(_node, UnaryOp::Sign) => nodes,
648
649 | Op::IfElse(_, None, None) => nodes,
650
651 | Op::BinaryScalarLhs(_, node, _)
652 | Op::BinaryScalarRhs(node, _, _)
653 | Op::Broadcast(node)
654 | Op::Unary(node, _)
655 | Op::Reduce(node, _, _)
656 | Op::Narrow(node, _, _, _)
657 | Op::Slice(node, _, _, _, _)
658 | Op::Reshape(node)
659 | Op::Transpose(node, _, _)
660 | Op::Permute(node, _)
661 | Op::Copy(node)
662 | Op::Gather(node, _, _)
663 | Op::IndexSelect(node, _, _)
664 | Op::IfElse(_, Some(node), None)
665 | Op::IfElse(_, None, Some(node)) => {
666 let (tg, nodes) = walk(node, nodes, already_seen);
667 track_grad |= tg;
668 nodes
669 }
670
671 | Op::Cat(args, _) => args.iter().fold(nodes, |nodes, arg| {
672 let (tg, nodes) = walk(arg, nodes, already_seen);
673 track_grad |= tg;
674 nodes
675 }),
676
677 | Op::RmsNorm(input, weight, _) => {
678 let (tg, nodes) = walk(input, nodes, already_seen);
679 track_grad |= tg;
680 let (tg, nodes) = walk(weight, nodes, already_seen);
681 track_grad |= tg;
682 nodes
683 }
684
685 | Op::Softmax(input, _) => {
686 let (tg, nodes) = walk(input, nodes, already_seen);
687 track_grad |= tg;
688 nodes
689 }
690 }
691 } else {
692 nodes
693 };
694 already_seen.insert(node.id(), track_grad);
695 if track_grad {
696 nodes.push(node);
697 }
698 (track_grad, nodes)
699 }
700 let (_tg, mut nodes) = walk(self, vec![], &mut HashMap::new());
701 nodes.reverse();
702 nodes
703 }
704
705 fn broadcast_back(arg: &Tensor<T>, node: &Tensor<T>, reduced_dims: &[usize]) -> crate::Result<Tensor<T>> {
706 if arg.rank() == node.rank() {
707 node.broadcast_as(arg.shape())
708 } else {
709 node.reshape(reduced_dims)?.broadcast_as(arg.shape())
710 }
711 }
712}