github-starred/cs249r_book
2
0
Fork 0
mirror of https://github.com/harvard-edge/cs249r_book.git synced 2026-07-16 14:42:29 -05:00
Code Issues 276 Packages Projects Releases 26 Wiki Activity

refactor: implement strict progressive disclosure for autograd

Browse Source 
- Module 01: Remove requires_grad, grad, backward() from Tensor class
  Students learn pure tensor math first without gradient concepts

- Module 02: Remove requires_grad propagation from Softmax
  Activations are now forward-only until autograd is enabled

- Module 03: Remove requires_grad=True from layer weights
  Layers store parameters without gradient flags

- Module 05: Update enable_autograd() to ADD gradient infrastructure
  Now adds requires_grad, grad, and backward() to Tensor class
  Uses helper functions for tensors created before autograd

- Module 09: Remove Conv2dBackward, MaxPool2dBackward classes
  Convolutions are now forward-only, no Module 05 import

- Module 11: Remove EmbeddingBackward import and usage
  Embeddings are now forward-only

- Modules 12, 13: Remove requires_grad from mask/param tensors

This implements true progressive disclosure: concepts are introduced
only when students are ready to learn them. Gradient tracking is now
completely absent from Modules 01-04 and added in Module 05.
This commit is contained in:
Vijay Janapa Reddi
2025-12-17 15:09:38 -05:00
parent 08ff168328
commit 23c5eb2b51
8 changed files with 156 additions and 425 deletions
Download Patch File Download Diff File Expand all files Collapse all files

159
tinytorch/src/01_tensor/01_tensor.py
View File

@@ -33,8 +33,8 @@ NumPy Arrays → Tensor → Activations (Module 02)
By the end of this module, you will:
1. Implement a complete Tensor class with fundamental operations
2. Understand tensors as the universal data structure in ML
3. Test tensor operations with immediate validation
4. Prepare for gradient computation in Module 05
3. Master broadcasting, matrix multiplication, and shape manipulation
4. Test tensor operations with immediate validation
Let's get started!
@@ -193,7 +193,7 @@ This memory layout affects performance in real ML workloads - algorithms that ac
Let's build our Tensor class step by step, testing each component as we go.
**Key Design Decision**: We'll include gradient-related attributes from the start, but they'll remain dormant until Module 05. This ensures a consistent interface throughout the course while keeping the cognitive load manageable.
**Key Design Decision**: We focus on core tensor operations first. Gradient tracking will be added in Module 05 (Autograd) when you're ready to learn backpropagation.
### Tensor Class Architecture
@@ -204,21 +204,24 @@ Tensor Class Structure:
│ • data: np.array (the numbers) │
│ • shape: tuple (dimensions) │
│ • size: int (total elements) │
│ • dtype: type (float32, int64) │
│ • dtype: type (float32) │
├─────────────────────────────────┤
│ Gradient Attributes (dormant): │
│ • requires_grad: bool │
│ • grad: None (until Module 05) │
├─────────────────────────────────┤
│ Operations: │
│ Arithmetic Operations: │
│ • __add__, __sub__, __mul__ │
│ • matmul(), reshape() │
│ • __truediv__, matmul() │
├─────────────────────────────────┤
│ Shape Operations: │
│ • reshape(), transpose() │
│ • sum(), mean(), max() │
│ • __getitem__ (indexing) │
├─────────────────────────────────┤
│ Utility Methods: │
│ • __repr__(), __str__() │
│ • numpy(), memory_footprint() │
└─────────────────────────────────┘
```
The beauty of this design: **all methods are defined inside the class from day one**. No monkey-patching, no dynamic attribute addition. Clean, consistent, debugger-friendly.
This clean design focuses on what tensors fundamentally do: store and manipulate numerical data efficiently.
"""
# %% [markdown]
@@ -255,31 +258,29 @@ Tensor wraps with: shape=(2,3), size=6, dtype=int64
# %% nbgrader={"grade": false, "grade_id": "tensor-class", "solution": true}
#| export
class Tensor:
"""Educational tensor that grows with student knowledge.
"""Educational tensor - the foundation of machine learning computation.
This class starts simple but includes dormant features for future modules:
- requires_grad: Will be used for automatic differentiation (Module 05)
- grad: Will store computed gradients (Module 05)
- backward(): Will compute gradients (Module 05)
This class provides the core data structure for all ML operations:
- data: The actual numerical values (NumPy array)
- shape: Dimensions of the tensor
- size: Total number of elements
- dtype: Data type (float32)
For now, focus on: data, shape, and basic operations.
All arithmetic, matrix, and shape operations are built on this foundation.
"""
def __init__(self, data, requires_grad=False):
def __init__(self, data):
"""Create a new tensor from data."""
### BEGIN SOLUTION
self.data = np.array(data, dtype=np.float32)
self.shape = self.data.shape
self.size = self.data.size
self.dtype = self.data.dtype
self.requires_grad = requires_grad
self.grad = None
### END SOLUTION
def __repr__(self):
"""String representation of tensor for debugging."""
grad_info = f", requires_grad={self.requires_grad}" if self.requires_grad else ""
return f"Tensor(data={self.data}, shape={self.shape}{grad_info})"
return f"Tensor(data={self.data}, shape={self.shape})"
def __str__(self):
"""Human-readable string representation."""
@@ -389,8 +390,7 @@ class Tensor:
result_data = self.data[key]
if not isinstance(result_data, np.ndarray):
result_data = np.array(result_data)
result = Tensor(result_data, requires_grad=self.requires_grad)
return result
return Tensor(result_data)
### END SOLUTION
def reshape(self, *shape):
@@ -418,8 +418,7 @@ class Tensor:
f"Total elements must match: {self.size} ≠ {target_size}"
)
reshaped_data = np.reshape(self.data, new_shape)
result = Tensor(reshaped_data, requires_grad=self.requires_grad)
return result
return Tensor(reshaped_data)
### END SOLUTION
def transpose(self, dim0=None, dim1=None):
@@ -438,8 +437,7 @@ class Tensor:
axes = list(range(len(self.shape)))
axes[dim0], axes[dim1] = axes[dim1], axes[dim0]
transposed_data = np.transpose(self.data, axes)
result = Tensor(transposed_data, requires_grad=self.requires_grad)
return result
return Tensor(transposed_data)
### END SOLUTION
def sum(self, axis=None, keepdims=False):
@@ -463,12 +461,6 @@ class Tensor:
return Tensor(result)
### END SOLUTION
def backward(self):
"""Compute gradients (implemented in Module 05: Autograd)."""
### BEGIN SOLUTION
pass
### END SOLUTION
# %% [markdown]
"""
### 🧪 Unit Test: Tensor Creation
@@ -490,8 +482,6 @@ def test_unit_tensor_creation():
assert scalar.data == 5.0
assert scalar.shape == ()
assert scalar.size == 1
assert scalar.requires_grad == False
assert scalar.grad is None
assert scalar.dtype == np.float32
# Test vector creation
@@ -506,10 +496,10 @@ def test_unit_tensor_creation():
assert matrix.shape == (2, 2)
assert matrix.size == 4
# Test gradient flag (dormant feature)
grad_tensor = Tensor([1, 2], requires_grad=True)
assert grad_tensor.requires_grad == True
assert grad_tensor.grad is None # Still None until Module 05
# Test 3D tensor creation
tensor_3d = Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
assert tensor_3d.shape == (2, 2, 2)
assert tensor_3d.size == 8
print("✅ Tensor creation works correctly!")
@@ -1128,76 +1118,6 @@ def test_unit_reduction_operations():
if __name__ == "__main__":
test_unit_reduction_operations()
# %% [markdown]
"""
## Gradient Features: Preparing for Module 05
Our Tensor includes dormant gradient features that will spring to life in Module 05. For now, they exist but do nothing - this design choice ensures a consistent interface throughout the course.
### Why Include Gradient Features Now?
```
Gradient System Evolution:
Module 01: Tensor with dormant gradients
┌─────────────────────────────────┐
│ Tensor │
│ • data: actual values │
│ • requires_grad: False │ ← Present but unused
│ • grad: None │ ← Present but stays None
│ • backward(): pass │ ← Present but does nothing
└─────────────────────────────────┘
↓ Module 05 activates these
Module 05: Tensor with active gradients
┌─────────────────────────────────┐
│ Tensor │
│ • data: actual values │
│ • requires_grad: True │ ← Now controls gradient tracking
│ • grad: computed gradients │ ← Now accumulates gradients
│ • backward(): computes grads │ ← Now implements chain rule
└─────────────────────────────────┘
```
### Design Benefits
**Consistency**: Same Tensor class interface throughout all modules
- No confusing Variable vs. Tensor distinction (unlike early PyTorch)
- Students never need to learn a "new" Tensor class
- IDE autocomplete works from day one
**Gradual Complexity**: Features activate when students are ready
- Module 01-04: Ignore gradient features, focus on operations
- Module 05: Gradient features "turn on" magically
- No cognitive overload in early modules
**Future-Proof**: Easy to extend without breaking changes
- Additional features can be added as dormant initially
- No monkey-patching or dynamic class modification
- Clean evolution path
### Current State (Module 01)
```
Gradient Features - Current Behavior:
┌─────────────────────────────────────────────────────────┐
│ Feature │ Current State │ Module 05 State │
├─────────────────────────────────────────────────────────┤
│ requires_grad │ False │ True (when needed) │
│ grad │ None │ np.array(...) │
│ backward() │ pass (no-op) │ Chain rule impl │
│ Operation chaining│ Not tracked │ Computation graph │
└─────────────────────────────────────────────────────────┘
Student Experience:
• Can call .backward() without errors (just does nothing)
• Can set requires_grad=True (just gets stored)
• Focus on understanding tensor operations first
• Gradients remain "mysterious" until Module 05 reveals them
```
This approach matches the pedagogical principle of "progressive disclosure" - reveal complexity only when students are ready to handle it.
"""
# %% [markdown]
"""
## Systems Analysis: Memory Layout and Performance
@@ -1390,18 +1310,6 @@ def test_module():
print("✅ Two-layer neural network computation works!")
# Test gradient attributes are preserved and functional
print("🧪 Integration Test: Gradient System Readiness...")
grad_tensor = Tensor([1, 2, 3], requires_grad=True)
result = grad_tensor + 5
assert grad_tensor.requires_grad == True, "requires_grad not preserved"
assert grad_tensor.grad is None, "grad should still be None"
# Test backward() doesn't crash (even though it does nothing)
grad_tensor.backward() # Should not raise any exception
print("✅ Gradient system ready for Module 05!")
# Test complex shape manipulations
print("🧪 Integration Test: Complex Shape Operations...")
data = Tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
@@ -1624,19 +1532,18 @@ Congratulations! You've built the foundational Tensor class that powers all mach
### Key Accomplishments
- **Built a complete Tensor class** with arithmetic operations, matrix multiplication, and shape manipulation
- **Implemented broadcasting semantics** that match NumPy for automatic shape alignment
- **Created dormant gradient features** that will activate in Module 05 (autograd)
- **Created reduction operations** (sum, mean, max) for loss computation and pooling
- **Added comprehensive ASCII diagrams** showing tensor operations visually
- **All methods defined INSIDE the class** (no monkey-patching) for clean, maintainable code
- **All tests pass ✅** (validated by `test_module()`)
### Systems Insights Discovered
- **Memory scaling**: Matrix operations create new tensors (3× memory during computation)
- **Broadcasting efficiency**: NumPy's automatic shape alignment vs. explicit operations
- **Cache behavior**: Row-wise access is faster than column-wise due to memory layout
- **Shape validation trade-offs**: Clear errors vs. performance in tight loops
- **Architecture decisions**: Dormant features vs. inheritance for clean evolution
### Ready for Next Steps
Your Tensor implementation enables all future modules! The dormant gradient features will spring to life in Module 05, and every neural network component will build on this foundation.
Your Tensor implementation enables all future modules! In Module 05 (Autograd), gradient tracking will be added to enable training. Every neural network component will build on this foundation.
Export with: `tito module complete 01_tensor`

19
tinytorch/src/02_activations/02_activations.py
View File

@@ -762,19 +762,18 @@ class Softmax:
"""
### BEGIN SOLUTION
# Numerical stability: subtract max to prevent overflow
# Use Tensor operations to preserve gradient flow!
x_max_data = np.max(x.data, axis=dim, keepdims=True)
x_max = Tensor(x_max_data, requires_grad=False) # max is not differentiable in this context
x_shifted = x - x_max # Tensor subtraction!
x_max = Tensor(x_max_data)
x_shifted = x - x_max # Tensor subtraction
# Compute exponentials (NumPy operation, but wrapped in Tensor)
exp_values = Tensor(np.exp(x_shifted.data), requires_grad=x_shifted.requires_grad)
# Compute exponentials
exp_values = Tensor(np.exp(x_shifted.data))
# Sum along dimension (Tensor operation)
# Sum along dimension
exp_sum_data = np.sum(exp_values.data, axis=dim, keepdims=True)
exp_sum = Tensor(exp_sum_data, requires_grad=exp_values.requires_grad)
exp_sum = Tensor(exp_sum_data)
# Normalize to get probabilities (Tensor division!)
# Normalize to get probabilities
result = exp_values / exp_sum
return result
### END SOLUTION
@@ -783,10 +782,6 @@ class Softmax:
"""Allows the activation to be called like a function."""
return self.forward(x, dim)
def backward(self, grad: Tensor) -> Tensor:
"""Compute gradient (implemented in Module 05)."""
pass # Will implement backward pass in Module 05
# %% [markdown]
"""
### 🔬 Unit Test: Softmax

65
tinytorch/src/03_layers/03_layers.py
View File

@@ -165,9 +165,9 @@ Let's build our layer system step by step. We'll implement two essential layer t
### Key Design Principles:
- All methods defined INSIDE classes (no monkey-patching)
- Parameter tensors have requires_grad=True (ready for Module 05)
- Forward methods return new tensors, preserving immutability
- parameters() method enables optimizer integration
- Gradient tracking will be added in Module 05 (Autograd)
"""
# %% [markdown]
@@ -211,7 +211,7 @@ class Layer:
Return list of trainable parameters.
Returns:
List of Tensor objects with requires_grad=True
List of Tensor objects (weights and biases)
"""
return [] # Base class has no parameters
@@ -282,7 +282,7 @@ class Linear(Layer):
APPROACH:
1. Create weight matrix (in_features, out_features) with Xavier scaling
2. Create bias vector (out_features,) initialized to zeros if bias=True
3. Set requires_grad=True for parameters (ready for Module 05)
3. Store as Tensor objects for use in forward pass
EXAMPLE:
>>> layer = Linear(784, 10) # MNIST classifier final layer
@@ -303,12 +303,12 @@ class Linear(Layer):
# Xavier/Glorot initialization for stable gradients
scale = np.sqrt(XAVIER_SCALE_FACTOR / in_features)
weight_data = np.random.randn(in_features, out_features) * scale
self.weight = Tensor(weight_data, requires_grad=True)
self.weight = Tensor(weight_data)
# Initialize bias to zeros or None
if bias:
bias_data = np.zeros(out_features)
self.bias = Tensor(bias_data, requires_grad=True)
self.bias = Tensor(bias_data)
else:
self.bias = None
### END SOLUTION
@@ -390,8 +390,6 @@ def test_unit_linear_layer():
assert layer.out_features == 256
assert layer.weight.shape == (784, 256)
assert layer.bias.shape == (256,)
assert layer.weight.requires_grad == True
assert layer.bias.requires_grad == True
# Test Xavier initialization (weights should be reasonably scaled)
weight_std = np.std(layer.weight.data)
@@ -467,35 +465,32 @@ if __name__ == "__main__":
# %% [markdown]
"""
### 🔬 Gradient Preparation Tests: Linear Layer
Tests to ensure Linear layer is ready for gradient-based training (Module 05).
### 🔬 Parameter Collection Tests: Linear Layer
Tests to ensure Linear layer parameters can be collected for optimization.
"""
# %% nbgrader={"grade": true, "grade_id": "test-linear-grad-prep", "locked": true, "points": 5}
def test_gradient_preparation_linear():
"""🔬 Test Linear layer is ready for gradients (Module 05)."""
print("🔬 Gradient Preparation Test: Linear Layer...")
# %% nbgrader={"grade": true, "grade_id": "test-linear-params", "locked": true, "points": 5}
def test_parameter_collection_linear():
"""🔬 Test Linear layer parameter collection."""
print("🔬 Parameter Collection Test: Linear Layer...")
layer = Linear(10, 5)
# Verify requires_grad is set
assert layer.weight.requires_grad == True, "Weight should require gradients"
assert layer.bias.requires_grad == True, "Bias should require gradients"
# Verify gradient placeholders exist (even if None initially)
assert hasattr(layer.weight, 'grad'), "Weight should have grad attribute"
assert hasattr(layer.bias, 'grad'), "Bias should have grad attribute"
# Verify parameter collection works
params = layer.parameters()
assert len(params) == 2, "Should return 2 parameters"
assert all(p.requires_grad for p in params), "All parameters should require gradients"
assert len(params) == 2, "Should return 2 parameters (weight and bias)"
assert params[0].shape == (10, 5), "First param should be weight"
assert params[1].shape == (5,), "Second param should be bias"
print("✅ Layer ready for gradient-based training!")
# Test layer without bias
layer_no_bias = Linear(10, 5, bias=False)
params_no_bias = layer_no_bias.parameters()
assert len(params_no_bias) == 1, "Should return 1 parameter (weight only)"
print("✅ Parameter collection works correctly!")
if __name__ == "__main__":
test_gradient_preparation_linear()
test_parameter_collection_linear()
# %% [markdown]
@@ -596,7 +591,7 @@ class Dropout(Layer):
APPROACH:
1. If training=False or p=0, return input unchanged
2. If p=1, return zeros (preserve requires_grad)
2. If p=1, return zeros
3. Otherwise: create random mask, apply it, scale by 1/(1-p)
EXAMPLE:
@@ -616,8 +611,8 @@ class Dropout(Layer):
return x
if self.p == DROPOUT_MAX_PROB:
# Drop everything (preserve requires_grad for gradient flow)
return Tensor(np.zeros_like(x.data), requires_grad=x.requires_grad)
# Drop everything
return Tensor(np.zeros_like(x.data))
# During training, apply dropout
keep_prob = 1.0 - self.p
@@ -625,9 +620,9 @@ class Dropout(Layer):
# Create random mask: True where we keep elements
mask = np.random.random(x.data.shape) < keep_prob
# Apply mask and scale using Tensor operations to preserve gradients!
mask_tensor = Tensor(mask.astype(np.float32), requires_grad=False) # Mask doesn't need gradients
scale = Tensor(np.array(1.0 / keep_prob), requires_grad=False)
# Apply mask and scale
mask_tensor = Tensor(mask.astype(np.float32))
scale = Tensor(np.array(1.0 / keep_prob))
# Use Tensor operations: x * mask * scale
output = x * mask_tensor * scale
@@ -1015,7 +1010,7 @@ def test_module():
print("Running unit tests...")
test_unit_linear_layer()
test_edge_cases_linear()
test_gradient_preparation_linear()
test_parameter_collection_linear()
test_unit_dropout_layer()
print("\nRunning integration scenarios...")
@@ -1055,10 +1050,6 @@ def test_module():
expected_params = 6 # 3 weights + 3 biases from 3 Linear layers
assert len(all_params) == expected_params, f"Expected {expected_params} parameters, got {len(all_params)}"
# Test all parameters have requires_grad=True
for param in all_params:
assert param.requires_grad == True, "All parameters should have requires_grad=True"
# Test individual layer functionality
test_x = Tensor(np.random.randn(4, 784))
# Test dropout in training vs inference

80
tinytorch/src/05_autograd/05_autograd.py
View File

@@ -1366,6 +1366,19 @@ def enable_autograd(quiet=False):
# Silently return if already enabled - no need to warn
return
# ===== STEP 1: Add gradient infrastructure to Tensor =====
# Store original __init__ to extend it
_original_init = Tensor.__init__
def gradient_aware_init(self, data, requires_grad=False):
"""Extended Tensor init that supports gradient tracking."""
_original_init(self, data)
self.requires_grad = requires_grad
self.grad = None
# Replace __init__ with gradient-aware version
Tensor.__init__ = gradient_aware_init
# Store original operations
# These are guaranteed to exist from Module 01 (Tensor class)
_original_add = Tensor.__add__
@@ -1379,6 +1392,19 @@ def enable_autograd(quiet=False):
_original_transpose = Tensor.transpose
_original_reshape = Tensor.reshape
# Helper to safely check requires_grad (handles tensors created before enable_autograd)
def _get_requires_grad(tensor):
"""Safely get requires_grad, defaulting to False for pre-autograd tensors."""
return getattr(tensor, 'requires_grad', False) if isinstance(tensor, Tensor) else False
def _ensure_grad_attrs(tensor):
"""Ensure tensor has gradient attributes (for tensors created before enable_autograd)."""
if isinstance(tensor, Tensor):
if not hasattr(tensor, 'requires_grad'):
tensor.requires_grad = False
if not hasattr(tensor, 'grad'):
tensor.grad = None
# Enhanced operations that track gradients
def tracked_add(self, other):
"""
@@ -1387,15 +1413,20 @@ def enable_autograd(quiet=False):
Enhances the original __add__ method to build computation graphs
when requires_grad=True for any input.
"""
# Ensure self has gradient attributes
_ensure_grad_attrs(self)
# Convert scalar to Tensor if needed
if not isinstance(other, Tensor):
other = Tensor(other)
_ensure_grad_attrs(other)
# Call original operation
result = _original_add(self, other)
_ensure_grad_attrs(result)
# Track gradient if needed
if self.requires_grad or other.requires_grad:
if _get_requires_grad(self) or _get_requires_grad(other):
result.requires_grad = True
result._grad_fn = AddBackward(self, other)
@@ -1408,17 +1439,21 @@ def enable_autograd(quiet=False):
Enhances the original __mul__ method to build computation graphs
when requires_grad=True for any input.
"""
_ensure_grad_attrs(self)
# Convert scalar to Tensor if needed for consistency
if not isinstance(other, Tensor):
other_tensor = Tensor(other)
else:
other_tensor = other
_ensure_grad_attrs(other_tensor)
# Call original operation
result = _original_mul(self, other)
_ensure_grad_attrs(result)
# Track gradient if needed
if self.requires_grad or (isinstance(other, Tensor) and other.requires_grad):
if _get_requires_grad(self) or _get_requires_grad(other_tensor):
result.requires_grad = True
result._grad_fn = MulBackward(self, other)
@@ -1431,11 +1466,15 @@ def enable_autograd(quiet=False):
Enhances the original matmul method to build computation graphs
when requires_grad=True for any input.
"""
_ensure_grad_attrs(self)
_ensure_grad_attrs(other)
# Call original matmul from Module 01
result = _original_matmul(self, other)
_ensure_grad_attrs(result)
# Track gradient if needed
if self.requires_grad or other.requires_grad:
if _get_requires_grad(self) or _get_requires_grad(other):
result.requires_grad = True
result._grad_fn = MatmulBackward(self, other)
@@ -1448,11 +1487,14 @@ def enable_autograd(quiet=False):
Enhances the original transpose method to build computation graphs
when requires_grad=True for the input.
"""
_ensure_grad_attrs(self)
# Call original transpose from Module 01
result = _original_transpose(self, dim0, dim1)
_ensure_grad_attrs(result)
# Track gradient if needed
if self.requires_grad:
if _get_requires_grad(self):
result.requires_grad = True
result._grad_fn = TransposeBackward(self, dim0, dim1)
@@ -1465,13 +1507,15 @@ def enable_autograd(quiet=False):
Enhances the original reshape method to build computation graphs
when requires_grad=True for the input.
"""
_ensure_grad_attrs(self)
original_shape = self.shape
# Call original reshape from Module 01
result = _original_reshape(self, *shape)
_ensure_grad_attrs(result)
# Track gradient if needed
if self.requires_grad:
if _get_requires_grad(self):
result.requires_grad = True
result._grad_fn = ReshapeBackward(self, original_shape)
@@ -1484,15 +1528,19 @@ def enable_autograd(quiet=False):
Enhances the original __sub__ method to build computation graphs
when requires_grad=True for any input.
"""
_ensure_grad_attrs(self)
# Convert scalar to Tensor if needed
if not isinstance(other, Tensor):
other = Tensor(other)
_ensure_grad_attrs(other)
# Call original operation
result = _original_sub(self, other)
_ensure_grad_attrs(result)
# Track gradient if needed
if self.requires_grad or other.requires_grad:
if _get_requires_grad(self) or _get_requires_grad(other):
result.requires_grad = True
result._grad_fn = SubBackward(self, other)
@@ -1505,15 +1553,19 @@ def enable_autograd(quiet=False):
Enhances the original __truediv__ method to build computation graphs
when requires_grad=True for any input.
"""
_ensure_grad_attrs(self)
# Convert scalar to Tensor if needed
if not isinstance(other, Tensor):
other = Tensor(other)
_ensure_grad_attrs(other)
# Call original operation
result = _original_div(self, other)
_ensure_grad_attrs(result)
# Track gradient if needed
if self.requires_grad or other.requires_grad:
if _get_requires_grad(self) or _get_requires_grad(other):
result.requires_grad = True
result._grad_fn = DivBackward(self, other)
@@ -1526,11 +1578,14 @@ def enable_autograd(quiet=False):
Enhances the original __getitem__ method to build computation graphs
when requires_grad=True for the input.
"""
_ensure_grad_attrs(self)
# Call original __getitem__ from Module 01
result = _original_getitem(self, key)
_ensure_grad_attrs(result)
# Track gradient if needed
if self.requires_grad:
if _get_requires_grad(self):
result.requires_grad = True
result._grad_fn = SliceBackward(self, key)
@@ -1543,10 +1598,12 @@ def enable_autograd(quiet=False):
Creates a new sum method that builds computation graphs
when requires_grad=True.
"""
_ensure_grad_attrs(self)
result_data = np.sum(self.data, axis=axis, keepdims=keepdims)
result = Tensor(result_data)
if self.requires_grad:
if _get_requires_grad(self):
result.requires_grad = True
result._grad_fn = SumBackward(self)
@@ -1573,8 +1630,11 @@ def enable_autograd(quiet=False):
print(x.grad) # [3.0]
```
"""
# Ensure gradient attributes exist
_ensure_grad_attrs(self)
# Only compute gradients if required
if not self.requires_grad:
if not _get_requires_grad(self):
return
# Initialize gradient if not provided (for scalar outputs)

212
tinytorch/src/09_convolutions/09_convolutions.py
View File

@@ -65,7 +65,6 @@ import numpy as np
import time
from tinytorch.core.tensor import Tensor
from tinytorch.core.autograd import Function
# Constants for convolution defaults
DEFAULT_KERNEL_SIZE = 3 # Default kernel size for convolutions
@@ -298,110 +297,6 @@ This reveals why convolution is expensive: O(B×C_out×H×W×K_h×K_w×C_in) ope
#| export
class Conv2dBackward(Function):
"""
Gradient computation for 2D convolution.
Computes gradients for Conv2d backward pass:
- grad_input: gradient w.r.t. input (for backprop to previous layer)
- grad_weight: gradient w.r.t. filters (for weight updates)
- grad_bias: gradient w.r.t. bias (for bias updates)
This uses explicit loops to show the gradient computation, matching
the educational approach of the forward pass.
"""
def __init__(self, x, weight, bias, stride, padding, kernel_size, padded_shape):
# Register all tensors that need gradients with autograd
if bias is not None:
super().__init__(x, weight, bias)
else:
super().__init__(x, weight)
self.x = x
self.weight = weight
self.bias = bias
self.stride = stride
self.padding = padding
self.kernel_size = kernel_size
self.padded_shape = padded_shape
def apply(self, grad_output):
"""
Compute gradients for convolution inputs and parameters.
Args:
grad_output: Gradient flowing back from next layer
Shape: (batch_size, out_channels, out_height, out_width)
Returns:
Tuple of (grad_input, grad_weight, grad_bias)
"""
batch_size, out_channels, out_height, out_width = grad_output.shape
_, in_channels, in_height, in_width = self.x.shape
kernel_h, kernel_w = self.kernel_size
# Apply padding to input if needed (for gradient computation)
if self.padding > 0:
padded_input = np.pad(self.x.data,
((0, 0), (0, 0), (self.padding, self.padding), (self.padding, self.padding)),
mode='constant', constant_values=0)
else:
padded_input = self.x.data
# Initialize gradients
grad_input_padded = np.zeros_like(padded_input)
grad_weight = np.zeros_like(self.weight.data)
grad_bias = None if self.bias is None else np.zeros_like(self.bias.data)
# Compute gradients using explicit loops (educational approach)
for b in range(batch_size):
for out_ch in range(out_channels):
for out_h in range(out_height):
for out_w in range(out_width):
# Position in input
in_h_start = out_h * self.stride
in_w_start = out_w * self.stride
# Gradient value flowing back to this position
grad_val = grad_output[b, out_ch, out_h, out_w]
# Distribute gradient to weight and input
for k_h in range(kernel_h):
for k_w in range(kernel_w):
for in_ch in range(in_channels):
# Input position
in_h = in_h_start + k_h
in_w = in_w_start + k_w
# Gradient w.r.t. weight
grad_weight[out_ch, in_ch, k_h, k_w] += (
padded_input[b, in_ch, in_h, in_w] * grad_val
)
# Gradient w.r.t. input
grad_input_padded[b, in_ch, in_h, in_w] += (
self.weight.data[out_ch, in_ch, k_h, k_w] * grad_val
)
# Compute gradient w.r.t. bias (sum over batch and spatial dimensions)
if grad_bias is not None:
for out_ch in range(out_channels):
grad_bias[out_ch] = grad_output[:, out_ch, :, :].sum()
# Remove padding from input gradient
if self.padding > 0:
grad_input = grad_input_padded[:, :,
self.padding:-self.padding,
self.padding:-self.padding]
else:
grad_input = grad_input_padded
# Return gradients as numpy arrays (autograd system handles storage)
# Following TinyTorch protocol: return (grad_input, grad_weight, grad_bias)
return grad_input, grad_weight, grad_bias
#| export
class Conv2d:
"""
2D Convolution layer for spatial feature extraction.
@@ -458,12 +353,11 @@ class Conv2d:
# Weight shape: (out_channels, in_channels, kernel_h, kernel_w)
self.weight = Tensor(np.random.normal(0, std,
(out_channels, in_channels, kernel_h, kernel_w)),
requires_grad=True)
(out_channels, in_channels, kernel_h, kernel_w)))
# Bias initialization
if bias:
self.bias = Tensor(np.zeros(out_channels), requires_grad=True)
self.bias = Tensor(np.zeros(out_channels))
else:
self.bias = None
### END SOLUTION
@@ -558,18 +452,7 @@ class Conv2d:
for out_ch in range(out_channels):
output[:, out_ch, :, :] += self.bias.data[out_ch]
# Return Tensor with gradient tracking enabled
result = Tensor(output, requires_grad=(x.requires_grad or self.weight.requires_grad))
# Attach backward function for gradient computation (following TinyTorch protocol)
if result.requires_grad:
result._grad_fn = Conv2dBackward(
x, self.weight, self.bias,
self.stride, self.padding, self.kernel_size,
padded_input.shape
)
return result
return Tensor(output)
### END SOLUTION
def parameters(self):
@@ -799,84 +682,6 @@ For input (1, 64, 224, 224) with 2×2 pooling:
#| export
class MaxPool2dBackward(Function):
"""
Gradient computation for 2D max pooling.
Max pooling gradients flow only to the positions that were selected
as the maximum in the forward pass.
"""
def __init__(self, x, output_shape, kernel_size, stride, padding):
super().__init__(x)
self.x = x
self.output_shape = output_shape
self.kernel_size = kernel_size
self.stride = stride
self.padding = padding
# Store max positions for gradient routing
self.max_positions = {}
def apply(self, grad_output):
"""
Route gradients back to max positions.
Args:
grad_output: Gradient from next layer
Returns:
Gradient w.r.t. input
"""
batch_size, channels, in_height, in_width = self.x.shape
_, _, out_height, out_width = self.output_shape
kernel_h, kernel_w = self.kernel_size
# Apply padding if needed
if self.padding > 0:
padded_input = np.pad(self.x.data,
((0, 0), (0, 0), (self.padding, self.padding), (self.padding, self.padding)),
mode='constant', constant_values=-np.inf)
grad_input_padded = np.zeros_like(padded_input)
else:
padded_input = self.x.data
grad_input_padded = np.zeros_like(self.x.data)
# Route gradients to max positions
for b in range(batch_size):
for c in range(channels):
for out_h in range(out_height):
for out_w in range(out_width):
in_h_start = out_h * self.stride
in_w_start = out_w * self.stride
# Find max position in this window
max_val = -np.inf
max_h, max_w = 0, 0
for k_h in range(kernel_h):
for k_w in range(kernel_w):
in_h = in_h_start + k_h
in_w = in_w_start + k_w
val = padded_input[b, c, in_h, in_w]
if val > max_val:
max_val = val
max_h, max_w = in_h, in_w
# Route gradient to max position
grad_input_padded[b, c, max_h, max_w] += grad_output[b, c, out_h, out_w]
# Remove padding
if self.padding > 0:
grad_input = grad_input_padded[:, :,
self.padding:-self.padding,
self.padding:-self.padding]
else:
grad_input = grad_input_padded
# Return as tuple (following Function protocol)
return (grad_input,)
#| export
class MaxPool2d:
"""
2D Max Pooling layer for spatial dimension reduction.
@@ -1000,16 +805,7 @@ class MaxPool2d:
# Store result
output[b, c, out_h, out_w] = max_val
# Return Tensor with gradient tracking
result = Tensor(output, requires_grad=x.requires_grad)
# Attach backward function for gradient computation
if result.requires_grad:
result._grad_fn = MaxPool2dBackward(
x, output.shape, self.kernel_size, self.stride, self.padding
)
return result
return Tensor(output)
### END SOLUTION
def parameters(self):

32
tinytorch/src/11_embeddings/11_embeddings.py
View File

@@ -66,7 +66,6 @@ from typing import List, Optional, Tuple
# Import from previous modules - following dependency chain
from tinytorch.core.tensor import Tensor
from tinytorch.core.autograd import EmbeddingBackward
# Constants for memory calculations
BYTES_PER_FLOAT32 = 4 # Standard float32 size in bytes
@@ -278,8 +277,7 @@ class Embedding:
# Xavier initialization for better gradient flow
limit = math.sqrt(6.0 / (vocab_size + embed_dim))
self.weight = Tensor(
np.random.uniform(-limit, limit, (vocab_size, embed_dim)),
requires_grad=True
np.random.uniform(-limit, limit, (vocab_size, embed_dim))
)
def forward(self, indices: Tensor) -> Tensor:
@@ -303,14 +301,7 @@ class Embedding:
# This is equivalent to one-hot multiplication but much more efficient
embedded = self.weight.data[indices.data.astype(int)]
# Create result tensor with gradient tracking
result = Tensor(embedded, requires_grad=self.weight.requires_grad)
# Attach backward function for gradient computation (following TinyTorch protocol)
if result.requires_grad:
result._grad_fn = EmbeddingBackward(self.weight, indices)
return result
return Tensor(embedded)
def __call__(self, indices: Tensor) -> Tensor:
"""Allows the embedding to be called like a function."""
@@ -355,7 +346,7 @@ def test_unit_embedding():
# Test 4: Parameter access
params = embed.parameters()
assert all(p.requires_grad for p in params), "All parameters should require gradients"
assert len(params) == 1, "Should have 1 parameter"
print("✅ Embedding layer works correctly!")
@@ -451,8 +442,7 @@ class PositionalEncoding:
# Smaller initialization than token embeddings since these are additive
limit = math.sqrt(2.0 / embed_dim)
self.position_embeddings = Tensor(
np.random.uniform(-limit, limit, (max_seq_len, embed_dim)),
requires_grad=True
np.random.uniform(-limit, limit, (max_seq_len, embed_dim))
)
def forward(self, x: Tensor) -> Tensor:
@@ -481,19 +471,13 @@ class PositionalEncoding:
)
# Slice position embeddings for this sequence length using Tensor slicing
# This now preserves gradient flow (as of Module 01 update with __getitem__)
pos_embeddings = self.position_embeddings[:seq_len] # (seq_len, embed_dim) - gradients preserved!
pos_embeddings = self.position_embeddings[:seq_len] # (seq_len, embed_dim)
# Reshape to add batch dimension: (1, seq_len, embed_dim)
# Need to use .data for reshaping temporarily, then wrap in Tensor
pos_data = pos_embeddings.data[np.newaxis, :, :]
pos_embeddings_batched = Tensor(pos_data, requires_grad=pos_embeddings.requires_grad)
pos_embeddings_batched = Tensor(pos_data)
# Copy gradient function if it exists (to preserve backward connection)
if hasattr(pos_embeddings, '_grad_fn') and pos_embeddings._grad_fn is not None:
pos_embeddings_batched._grad_fn = pos_embeddings._grad_fn
# Add positional information - gradients flow through both x and pos_embeddings!
# Add positional information
result = x + pos_embeddings_batched
return result
@@ -931,7 +915,7 @@ class EmbeddingLayer:
# Reshape to add batch dimension
pos_data = pos_embeddings.data[np.newaxis, :, :]
pos_embeddings_batched = Tensor(pos_data, requires_grad=False) # Sinusoidal are fixed
pos_embeddings_batched = Tensor(pos_data) # Sinusoidal are fixed
output = token_embeds + pos_embeddings_batched
else:

4
tinytorch/src/12_attention/12_attention.py
View File

@@ -327,7 +327,7 @@ def scaled_dot_product_attention(Q: Tensor, K: Tensor, V: Tensor, mask: Optional
# Ensure mask is broadcastable
mask_data = mask.data
adder_mask = (1.0 - mask_data) * MASK_VALUE
adder_mask_tensor = Tensor(adder_mask, requires_grad=False)
adder_mask_tensor = Tensor(adder_mask)
scores = scores + adder_mask_tensor
# Step 5: Apply softmax to get attention weights
@@ -648,7 +648,7 @@ class MultiHeadAttention:
# This allows the mask to broadcast across all attention heads
batch_size_mask, seq_len_mask, _ = mask.shape
mask_data = mask.data.reshape(batch_size_mask, 1, seq_len_mask, seq_len_mask)
mask_reshaped = Tensor(mask_data, requires_grad=False)
mask_reshaped = Tensor(mask_data)
attended, _ = scaled_dot_product_attention(Q, K, V, mask=mask_reshaped)

10
tinytorch/src/13_transformers/13_transformers.py
View File

@@ -435,8 +435,8 @@ class LayerNorm:
self.eps = eps
# Learnable parameters: scale and shift
self.gamma = Tensor(np.ones(normalized_shape), requires_grad=True) # Scale parameter
self.beta = Tensor(np.zeros(normalized_shape), requires_grad=True) # Shift parameter
self.gamma = Tensor(np.ones(normalized_shape)) # Scale parameter
self.beta = Tensor(np.zeros(normalized_shape)) # Shift parameter
### END SOLUTION
def forward(self, x):
@@ -465,10 +465,8 @@ class LayerNorm:
diff = x - mean
variance = (diff * diff).mean(axis=-1, keepdims=True)
# Normalize - use Tensor operations to preserve gradients!
# Add eps as a Tensor for proper gradient flow
eps_tensor = Tensor(np.array(self.eps), requires_grad=False)
std = Tensor(np.sqrt(variance.data + self.eps), requires_grad=variance.requires_grad)
# Normalize
std = Tensor(np.sqrt(variance.data + self.eps))
normalized = (x - mean) / std
# Apply learnable transformation