# Autograd Engine

How automatic differentiation works in RustyTorch++.

## Overview

The autograd engine implements reverse-mode automatic differentiation (backpropagation) using a tape-based approach.

## Key Concepts

### Computation Graph

Every tensor operation builds a computation graph:

```
     ┌───────────────┐
     │   Input X     │  (leaf, requires_grad=true)
     └───────┬───────┘
             │
     ┌───────▼───────┐
     │   MatMul(W)   │  W is another leaf
     └───────┬───────┘
             │
     ┌───────▼───────┐
     │     ReLU      │
     └───────┬───────┘
             │
     ┌───────▼───────┐
     │     Loss      │  (output)
     └───────────────┘
```

### Gradient Tape

Operations are recorded on a tape:

```rust
use rtx_autograd::{AutogradTape, TapeNode};

// Tape records operations during forward pass
let tape = AutogradTape::new();

// Each operation adds a node
tape.record_operation(
    &[input_id, weight_id],  // inputs
    output_id,                // output
    Box::new(MatMulBackward { /* saved tensors */ }),
);
```

## Using Variables

### Basic Gradient Computation

```rust
use rtx_autograd::{Variable, backward};
use rtx_tensor::{Tensor, Device};

let device = Device::cpu();

// Create tracked variables
let x = Variable::from_tensor(Tensor::from_slice(&[2.0, 3.0], [2], &device)?);
let w = Variable::from_tensor(Tensor::from_slice(&[1.0, 2.0], [2], &device)?);

// Forward pass builds graph
let y = x.multiply(&w)?;  // [2.0, 6.0]
let z = y.sum()?;         // 8.0

// Backward pass computes gradients
z.backward()?;

// Access gradients
println!("dx: {:?}", x.grad());  // [1.0, 2.0] (dy/dx = w)
println!("dw: {:?}", w.grad());  // [2.0, 3.0] (dy/dw = x)
```

### Gradient Context Managers

```rust
use rtx_autograd::{no_grad, enable_grad, is_grad_enabled};

// Check current state
println!("Grad enabled: {}", is_grad_enabled());

// Disable gradients (inference mode)
let result = no_grad(|| {
    let a = Tensor::randn([100, 100], &device)?;
    let b = Tensor::randn([100, 100], &device)?;
    a.matmul(&b)  // No graph built
});

// Re-enable within nested scope
no_grad(|| {
    enable_grad(|| {
        // Gradients tracked here
    });
    // Gradients not tracked here
});
```

### Detaching from Graph

```rust
let x = Variable::from_tensor(tensor);
let y = x.multiply(&w)?;

// Detach creates a new variable without graph connection
let y_detached = y.detach();
assert!(!y_detached.requires_grad());

// Useful for stop-gradient operations
let target = predictions.detach();  // Don't backprop through target
```

## Backward Functions

Each operation has a corresponding backward function:

```rust
use rtx_autograd::{BackwardFunction, BackwardContext};

/// Backward function for matrix multiplication: C = A @ B
pub struct MatMulBackward {
    a: Tensor,  // Saved input A
    b: Tensor,  // Saved input B
}

impl BackwardFunction for MatMulBackward {
    fn backward(&self, grad_output: &Tensor) -> Vec<Option<Tensor>> {
        // dL/dA = dL/dC @ B^T
        let grad_a = grad_output.matmul(&self.b.transpose(0, 1)?)?;

        // dL/dB = A^T @ dL/dC
        let grad_b = self.a.transpose(0, 1)?.matmul(grad_output)?;

        vec![Some(grad_a), Some(grad_b)]
    }
}
```

### Built-in Backward Functions

| Operation | Backward | Notes |
|-----------|----------|-------|
| `Add` | [1, 1] | Gradients pass through unchanged |
| `Mul` | [y, x] | Swap inputs |
| `MatMul` | [grad @ B^T, A^T @ grad] | Standard matrix calculus |
| `Sum` | expand | Broadcast gradient |
| `ReLU` | grad * (x > 0) | Zero where input was negative |
| `Sigmoid` | grad * σ * (1 - σ) | Uses saved output |
| `Softmax` | Jacobian-vector product | Complex formula |

## Gradient Checkpointing

Trade memory for compute by recomputing activations:

```rust
use rtx_autograd::{CheckpointManager, CheckpointStrategy};

// Create checkpoint manager
let mut checkpoint = CheckpointManager::new(CheckpointStrategy::SqrtN);

// Checkpoint expensive operations
let hidden = checkpoint.checkpoint(|| {
    model.layer1.forward(&input)?
        .relu()?
        .layer2.forward()
})?;

// During backward, activations are recomputed
hidden.backward()?;  // Recomputes forward pass for gradients
```

### Checkpointing Strategies

| Strategy | Memory | Compute | Use Case |
|----------|--------|---------|----------|
| `None` | O(n) | O(1) | Small models |
| `SqrtN` | O(√n) | O(√n) | Large transformers |
| `Every(k)` | O(n/k) | O(k) | Custom balance |
| `Selective` | Varies | Varies | Manual selection |

## Higher-Order Derivatives

### Hessian Computation

```rust
use rtx_autograd::{HessianComputer};

let f = |x: &Variable| -> Variable {
    x.pow(2)?.sum()?  // f(x) = sum(x^2)
};

let x = Variable::from_tensor(Tensor::from_slice(&[1.0, 2.0, 3.0], [3], &device)?);

// Compute Hessian (second derivative matrix)
let hessian = HessianComputer::compute(&f, &x)?;
// For sum(x^2), Hessian is 2*I (diagonal of 2s)
```

### Jacobian Computation

```rust
use rtx_autograd::JacobianComputer;

let f = |x: &Variable| -> Variable {
    // f: R^3 -> R^2
    let a = x.narrow(0, 0, 2)?;
    let b = x.narrow(0, 1, 2)?;
    a.add(&b)?
};

let x = Variable::from_tensor(Tensor::randn([3], &device)?);
let jacobian = JacobianComputer::compute(&f, &x)?;  // [2, 3] matrix
```

## Numerical Gradient Checking

Verify gradient implementations with finite differences:

```rust
use rtx_autograd::GradientChecker;

let checker = GradientChecker::new(1e-5);  // epsilon for finite diff

let f = |x: &Tensor| {
    x.pow(3)?.sum()  // f(x) = sum(x^3)
};

let x = Tensor::randn([10], &device)?;
let analytical_grad = compute_analytical_grad(&f, &x)?;

// Check matches numerical gradient
let is_correct = checker.check(&f, &x, &analytical_grad)?;
assert!(is_correct, "Gradient implementation is incorrect");
```

## Graph Caching

Cache computation graphs for repeated forward passes:

```rust
use rtx_autograd::{init_graph_cache, cache_graph, get_cached_graph};

// Initialize graph cache
init_graph_cache(1024);  // Cache up to 1024 graphs

// First forward pass - builds and caches graph
let output = model.forward(&input)?;

// Subsequent passes with same shapes - reuse graph
let cached = get_cached_graph(&input_shape)?;
if let Some(graph) = cached {
    // Execute cached graph
    output = graph.execute(&input)?;
}
```

## Best Practices

### 1. Use no_grad for Inference

```rust
// Wrong - builds unnecessary graph
let output = model.forward(&input)?;

// Right - skips graph construction
let output = no_grad(|| model.forward(&input))?;
```

### 2. Zero Gradients Before Each Step

```rust
for epoch in 0..num_epochs {
    for (x, y) in &dataloader {
        // Clear accumulated gradients
        optimizer.zero_grad();

        let output = model.forward(&x)?;
        let loss = loss_fn(&output, &y)?;

        loss.backward()?;
        optimizer.step()?;
    }
}
```

### 3. Detach When Needed

```rust
// Stop gradient flow to targets
let target = teacher_model.forward(&x)?.detach();
let student_out = student_model.forward(&x)?;
let loss = mse_loss(&student_out, &target)?;
```

### 4. Use Checkpointing for Large Models

```rust
// Memory-efficient transformer
let attention = checkpoint.checkpoint(|| {
    self.attention.forward(&hidden)?
})?;
```

## Next Steps

- [GPU Runtime](./gpu-runtime.md) - Device and memory management
- [Performance](../performance/gpu-optimization.md) - Optimization
- [Troubleshooting](../troubleshooting/memory.md) - Memory issues