⚡ Bolt: Optimize autograd backpropagation pass

.jules/bolt.md

@@ -0,0 +1,3 @@
+## 2026-07-05 - Optimize Autograd Backpropagation
+**Learning:** In reverse-mode backpropagation, cloning `Tensor` metadata for intermediate gradients causes redundant heap allocations. Using `Option::take()` to acquire ownership of gradients during backprop and re-insertion avoids borrow checker issues and unnecessary cloning.
+**Action:** Pass `Tensor` by value in `accumulate_grad` to utilize empty `None` slots efficiently.

crates/aether-core/src/ml/autograd.rs

@@ -200,52 +200,52 @@ impl<'a> Context<'a> {
             match op {
                 Op::Add { out, lhs, rhs } => {
                     // Solves borrow checker by cloning Option first
-                    let grad_out = grads[out.index].clone();
-                    if let Some(grad) = grad_out {
+                    if let Some(grad) = grads[out.index].take() {
                         // dL/d(lhs) += dL/dout * 1
-                        Self::accumulate_grad(&mut grads, lhs.index, &grad);
-                        Self::accumulate_grad(&mut grads, rhs.index, &grad);
+                        Self::accumulate_grad(&mut grads, lhs.index, grad.clone());
+                        Self::accumulate_grad(&mut grads, rhs.index, grad.clone());
+                        grads[out.index] = Some(grad);
                     }
                 }
                 Op::Mul { out, lhs, rhs } => {
-                    let grad_out = grads[out.index].clone();
-                    if let Some(grad) = grad_out {
+                    if let Some(grad) = grads[out.index].take() {
                         let lhs_val = self.heap.get(*lhs).unwrap();
                         let rhs_val = self.heap.get(*rhs).unwrap();
 
                         // dL/dLhs = grad_out * rhs
                         let d_lhs: Tensor = grad.mul(rhs_val);
-                        Self::accumulate_grad(&mut grads, lhs.index, &d_lhs);
+                        Self::accumulate_grad(&mut grads, lhs.index, d_lhs);
 
                         // dL/dRhs = grad_out * lhs
                         let d_rhs: Tensor = grad.mul(lhs_val);
-                        Self::accumulate_grad(&mut grads, rhs.index, &d_rhs);
+                        Self::accumulate_grad(&mut grads, rhs.index, d_rhs);
+                        grads[out.index] = Some(grad);
                     }
                 }
                 Op::MatMul { out, lhs, rhs } => {
-                    let grad_out = grads[out.index].clone();
-                    if let Some(grad) = grad_out {
+                    if let Some(grad) = grads[out.index].take() {
                         let lhs_val = self.heap.get(*lhs).unwrap();
                         let rhs_val = self.heap.get(*rhs).unwrap();
 
                         // C = A @ B
                         // dA = dC @ B^T
                         let d_lhs: Tensor = grad.matmul(&rhs_val.transpose());
-                        Self::accumulate_grad(&mut grads, lhs.index, &d_lhs);
+                        Self::accumulate_grad(&mut grads, lhs.index, d_lhs);
 
                         // dB = A^T @ dC
                         let d_rhs: Tensor = lhs_val.transpose().matmul(&grad);
-                        Self::accumulate_grad(&mut grads, rhs.index, &d_rhs);
+                        Self::accumulate_grad(&mut grads, rhs.index, d_rhs);
+                        grads[out.index] = Some(grad);
                     }
                 }
                 Op::ReLU { out, input } => {
-                    let grad_out = grads[out.index].clone();
-                    if let Some(grad) = grad_out {
+                    if let Some(grad) = grads[out.index].take() {
                         let input_val = self.heap.get(*input).unwrap();
                         // dL/dx = grad_out * (1 if x > 0 else 0)
                         let mask = input_val.map(|x| if x > 0.0 { 1.0 } else { 0.0 });
                         let d_input: Tensor = grad.mul(&mask);
-                        Self::accumulate_grad(&mut grads, input.index, &d_input);
+                        Self::accumulate_grad(&mut grads, input.index, d_input);
+                        grads[out.index] = Some(grad);
                     }
                 }
             }
@@ -254,18 +254,18 @@ impl<'a> Context<'a> {
         grads
     }
 
-    fn accumulate_grad(grads: &mut Vec<Option<Tensor>>, idx: usize, grad: &Tensor) {
+    fn accumulate_grad(grads: &mut Vec<Option<Tensor>>, idx: usize, grad: Tensor) {
         if idx >= grads.len() {
             grads.resize(idx + 1 + 256, None);
         }
 
         match &mut grads[idx] {
             Some(existing) => {
-                let new = existing.add(grad);
+                let new = existing.add(&grad);
                 grads[idx] = Some(new);
             }
             None => {
-                grads[idx] = Some(grad.clone());
+                grads[idx] = Some(grad);
             }
         }
     }

test_autograd.sh

@@ -0,0 +1,2 @@
+cd crates/aether-core
+cargo test -p aether-core --offline -- autograd