1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

import numpy as np
 
# ===== VJP Examples for Common Operations =====
 
def add_vjp(v, x, y):
    """
    VJP for z = x + y
    
    Jacobian: J = [I, I] where I is identity
    VJP: v^T @ [I, I] = [v, v]
    
    Gradient flows through unchanged to both inputs.
    """
    grad_x = v  # ∂L/∂x = v * 1
    grad_y = v  # ∂L/∂y = v * 1
    return grad_x, grad_y
 
def multiply_vjp(v, x, y):
    """
    VJP for z = x * y (elementwise)
    
    ∂z/∂x = y, ∂z/∂y = x (elementwise)
    """
    grad_x = v * y  # ∂L/∂x = v * y
    grad_y = v * x  # ∂L/∂y = v * x
    return grad_x, grad_y
 
def matmul_vjp(v, X, W):
    """
    VJP for Y = X @ W
    
    X: (batch, n), W: (n, m), Y: (batch, m)
    v: (batch, m) - upstream gradient
    
    ∂L/∂X = v @ W^T  (shape: batch, n)
    ∂L/∂W = X^T @ v  (shape: n, m)
    """
    grad_X = v @ W.T
    grad_W = X.T @ v
    return grad_X, grad_W
 
def relu_vjp(v, x, output):
    """
    VJP for y = ReLU(x) = max(0, x)
    
    ∂y/∂x = 1 if x > 0 else 0
    """
    grad_x = v * (x > 0).astype(float)
    return grad_x
 
def softmax_cross_entropy_vjp(v, logits, labels):
    """
    VJP for L = -log(softmax(logits)[label])
    
    The beautiful result: gradient = softmax(logits) - one_hot(label)
    This is efficient and numerically stable!
    """
    # Compute softmax
    exp_logits = np.exp(logits - np.max(logits, axis=-1, keepdims=True))
    softmax = exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)
    
    # Gradient is softmax - one_hot (for each example in batch)
    batch_size = logits.shape[0]
    grad_logits = softmax.copy()
    grad_logits[np.arange(batch_size), labels] -= 1
    
    # Scale by upstream gradient
    grad_logits *= v.reshape(-1, 1)
    
    return grad_logits
 
# Demonstration
print("=== VJP for Y = X @ W ===")
X = np.array([[1, 2], [3, 4]], dtype=float)
W = np.array([[0.5, 0.3, 0.2], [0.1, 0.4, 0.5]], dtype=float)
Y = X @ W
 
# Suppose upstream gradient is all ones
v = np.ones_like(Y)
grad_X, grad_W = matmul_vjp(v, X, W)
 
print(f"X shape: {X.shape}, W shape: {W.shape}")
print(f"Y shape: {Y.shape}, v shape: {v.shape}")
print(f"grad_X shape: {grad_X.shape} (same as X)")
print(f"grad_W shape: {grad_W.shape} (same as W)")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

def backward_pass(output_node):
    """
    Compute gradients for all nodes via reverse mode autodiff.
    
    Assumes forward pass has been run and values are cached.
    """
    # Step 1: Initialize gradients
    grads = {node: 0 for node in all_nodes}
    grads[output_node] = 1.0  # ∂L/∂L = 1
    
    # Step 2: Get reverse topological order
    reverse_topo = reversed(topological_sort(graph))
    
    # Step 3: Propagate gradients backward
    for node in reverse_topo:
        if not node.inputs:
            continue  # Input node, no upstream
        
        # Get upstream gradient for this node
        upstream_grad = grads[node]
        
        if upstream_grad == 0:
            continue  # No gradient flows (optimization)
        
        # Get cached forward values
        input_values = [inp.cached_value for inp in node.inputs]
        output_value = node.cached_value
        
        # Compute VJP: gradients w.r.t each input
        input_grads = node.vjp_function(
            upstream_grad, output_value, *input_values
        )
        
        # Step 4: ACCUMULATE into input nodes (handle fan-out)
        for inp_node, inp_grad in zip(node.inputs, input_grads):
            grads[inp_node] += inp_grad  # += not = !!!
    
    return grads
 
# Example usage
class MultiplyNode:
    def __init__(self, x_node, y_node):
        self.inputs = [x_node, y_node]
        self.cached_value = None
    
    def forward(self):
        x, y = self.inputs[0].cached_value, self.inputs[1].cached_value
        self.cached_value = x * y
        return self.cached_value
    
    def vjp_function(self, upstream_grad, output_value, x, y):
        # z = x * y
        # ∂z/∂x = y, ∂z/∂y = x
        grad_x = upstream_grad * y
        grad_y = upstream_grad * x
        return [grad_x, grad_y]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183

import numpy as np
from typing import List, Tuple, Callable, Optional
 
class Tensor:
    """
    A tensor with automatic differentiation support.
    
    Tracks computation graph via _children and computes gradients via _backward.
    """
    
    def __init__(
        self, 
        data: np.ndarray, 
        requires_grad: bool = False,
        _children: Tuple['Tensor', ...] = (),
        _op: str = ''
    ):
        self.data = np.array(data, dtype=np.float64)
        self.requires_grad = requires_grad
        self.grad: Optional[np.ndarray] = None
        
        # Graph structure for autograd
        self._children = set(_children)
        self._op = _op
        self._backward: Callable[[], None] = lambda: None
        
    def __repr__(self) -> str:
        return f"Tensor({self.data}, grad={self.grad})"
    
    @property
    def shape(self) -> Tuple[int, ...]:
        return self.data.shape
    
    def zero_grad(self) -> None:
        """Zero out gradient."""
        self.grad = None
    
    # ============ Operations ============
    
    def __add__(self, other: 'Tensor') -> 'Tensor':
        other = other if isinstance(other, Tensor) else Tensor(other)
        out = Tensor(
            self.data + other.data,
            requires_grad=self.requires_grad or other.requires_grad,
            _children=(self, other),
            _op='add'
        )
        
        def _backward():
            if self.requires_grad:
                # Handle broadcasting: sum over broadcasted dimensions
                grad = out.grad
                # Reduce grad to match self.shape if broadcasting occurred
                for i, (s1, s2) in enumerate(zip(out.shape, self.shape + (1,) * (len(out.shape) - len(self.shape)))):
                    if s1 != s2:
                        grad = grad.sum(axis=i, keepdims=True)
                grad = grad.reshape(self.shape) if grad.shape != self.shape else grad
                self.grad = (self.grad or 0) + grad
                
            if other.requires_grad:
                grad = out.grad
                for i, (s1, s2) in enumerate(zip(out.shape, other.shape + (1,) * (len(out.shape) - len(other.shape)))):
                    if s1 != s2:
                        grad = grad.sum(axis=i, keepdims=True)
                grad = grad.reshape(other.shape) if grad.shape != other.shape else grad
                other.grad = (other.grad or 0) + grad
                
        out._backward = _backward
        return out
    
    def __mul__(self, other: 'Tensor') -> 'Tensor':
        other = other if isinstance(other, Tensor) else Tensor(other)
        out = Tensor(
            self.data * other.data,
            requires_grad=self.requires_grad or other.requires_grad,
            _children=(self, other),
            _op='mul'
        )
        
        def _backward():
            if self.requires_grad:
                grad = out.grad * other.data
                # Handle broadcasting
                while len(grad.shape) > len(self.shape):
                    grad = grad.sum(axis=0)
                for i, (d1, d2) in enumerate(zip(self.shape, grad.shape)):
                    if d1 == 1 and d2 > 1:
                        grad = grad.sum(axis=i, keepdims=True)
                self.grad = (self.grad or 0) + grad
                
            if other.requires_grad:
                grad = out.grad * self.data
                while len(grad.shape) > len(other.shape):
                    grad = grad.sum(axis=0)
                for i, (d1, d2) in enumerate(zip(other.shape, grad.shape)):
                    if d1 == 1 and d2 > 1:
                        grad = grad.sum(axis=i, keepdims=True)
                other.grad = (other.grad or 0) + grad
                
        out._backward = _backward
        return out
    
    def __neg__(self) -> 'Tensor':
        return self * Tensor(-1.0)
    
    def __sub__(self, other: 'Tensor') -> 'Tensor':
        return self + (-other)
    
    def relu(self) -> 'Tensor':
        out = Tensor(
            np.maximum(0, self.data),
            requires_grad=self.requires_grad,
            _children=(self,),
            _op='relu'
        )
        
        def _backward():
            if self.requires_grad:
                self.grad = (self.grad or 0) + (self.data > 0) * out.grad
                
        out._backward = _backward
        return out
    
    def sum(self) -> 'Tensor':
        out = Tensor(
            self.data.sum(),
            requires_grad=self.requires_grad,
            _children=(self,),
            _op='sum'
        )
        
        def _backward():
            if self.requires_grad:
                # Gradient broadcasts to all elements
                self.grad = (self.grad or 0) + np.ones_like(self.data) * out.grad
                
        out._backward = _backward
        return out
    
    def mean(self) -> 'Tensor':
        out = Tensor(
            self.data.mean(),
            requires_grad=self.requires_grad,
            _children=(self,),
            _op='mean'
        )
        
        def _backward():
            if self.requires_grad:
                n = self.data.size
                self.grad = (self.grad or 0) + np.ones_like(self.data) * out.grad / n
                
        out._backward = _backward
        return out
    
    # ============ Backward Pass ============
    
    def backward(self) -> None:
        """
        Compute gradients via reverse-mode autodiff.
        Must be called on a scalar tensor.
        """
        assert self.data.size == 1, "backward() requires scalar output"
        
        # Build topological order
        topo: List['Tensor'] = []
        visited = set()
        
        def build_topo(tensor: 'Tensor'):
            if tensor not in visited:
                visited.add(tensor)
                for child in tensor._children:
                    build_topo(child)
                topo.append(tensor)
        
        build_topo(self)
        
        # Initialize gradient of output
        self.grad = np.ones_like(self.data)
        
        # Backpropagate in reverse topological order
        for tensor in reversed(topo):
            tensor._backward()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

# ============ Using Our Autograd System ============
 
# Example 1: Simple expression
print("=== Example 1: f(x,y) = (x+y)·x ===")
x = Tensor([3.0], requires_grad=True)
y = Tensor([2.0], requires_grad=True)
 
a = x + y  # a = 5
L = a * x  # L = 15
 
L.backward()
print(f"L = {L.data[0]}")
print(f"∂L/∂x = {x.grad[0]} (expected: 2x+y = 8)")
print(f"∂L/∂y = {y.grad[0]} (expected: x = 3)")
 
# Example 2: Neural network layer
print("
=== Example 2: Linear layer + ReLU ===")
# y = relu(x @ W + b)
np.random.seed(42)
x = Tensor(np.random.randn(4, 3), requires_grad=False)  # Input (batch=4, features=3)
W = Tensor(np.random.randn(3, 2) * 0.1, requires_grad=True)  # Weights
b = Tensor(np.zeros(2), requires_grad=True)  # Bias
 
# Forward pass
z = Tensor(x.data @ W.data) + b  # Linear combination
h = z.relu()                      # Activation
loss = h.mean()                   # Scalar loss
 
print(f"Output shape: {h.shape}")
print(f"Loss: {loss.data}")
 
# Backward pass
loss.backward()
 
print(f"W gradient shape: {W.grad.shape} (same as W)")
print(f"b gradient shape: {b.grad.shape} (same as b)")
print(f"W gradient:\n{W.grad}")
 
# Example 3: Verify correctness with numerical gradients
print("
=== Gradient Check ===")
def numerical_gradient(f, x, eps=1e-5):
    grad = np.zeros_like(x.data)
    for idx in np.ndindex(x.data.shape):
        x.data[idx] += eps
        loss_plus = f().data
        x.data[idx] -= 2 * eps
        loss_minus = f().data
        x.data[idx] += eps  # restore
        grad[idx] = (loss_plus - loss_minus) / (2 * eps)
    return grad
 
# Test gradient for a simple function
w = Tensor([2.0, 3.0], requires_grad=True)
 
def compute_loss():
    return ((w * w) + w * Tensor([1.0, 1.0])).sum()
 
loss = compute_loss()
loss.backward()
 
numerical_grad = numerical_gradient(compute_loss, w)
 
print(f"Autograd gradient: {w.grad}")
print(f"Numerical gradient: {numerical_grad}")
print(f"Difference: {np.abs(w.grad - numerical_grad).max():.2e}")