mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-07-20 07:58:34 -05:00
- Update tensor module notebook - Enhance activations module - Expand layers module functionality - Improve autograd implementation - Add optimizers enhancements - Update training module - Refine dataloader notebook
96 KiB
96 KiB
In [ ]:
#| default_exp core.tensor
#| export
import numpy as npIn [ ]:
#| export
class Tensor:
"""Educational tensor that grows with student knowledge.
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)
For now, focus on: data, shape, and basic operations.
"""
def __init__(self, data, requires_grad=False):
"""
Create a new tensor from data.
TODO: Initialize tensor attributes
APPROACH:
1. Convert data to NumPy array - handles lists, scalars, etc.
2. Store shape and size for quick access
3. Set up gradient tracking (dormant until Module 05)
EXAMPLE:
>>> tensor = Tensor([1, 2, 3])
>>> print(tensor.data)
[1 2 3]
>>> print(tensor.shape)
(3,)
HINT: np.array() handles type conversion automatically
"""
### BEGIN SOLUTION
# Core tensor data - always present
self.data = np.array(data, dtype=np.float32) # Consistent float32 for ML
self.shape = self.data.shape
self.size = self.data.size
self.dtype = self.data.dtype
# Gradient features (dormant until Module 05)
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})"
def __str__(self):
"""Human-readable string representation."""
return f"Tensor({self.data})"
def numpy(self):
"""Return the underlying NumPy array."""
return self.data
# nbgrader={\"grade\": false, \"grade_id\": \"addition-impl\", \"solution\": true}
def __add__(self, other):
"""
Add two tensors element-wise with broadcasting support.
TODO: Implement tensor addition with automatic broadcasting
APPROACH:
1. Handle both Tensor and scalar inputs
2. Use NumPy's broadcasting for automatic shape alignment
3. Return new Tensor with result (don't modify self)
EXAMPLE:
>>> a = Tensor([1, 2, 3])
>>> b = Tensor([4, 5, 6])
>>> result = a + b
>>> print(result.data)
[5. 7. 9.]
BROADCASTING EXAMPLE:
>>> matrix = Tensor([[1, 2], [3, 4]]) # Shape: (2, 2)
>>> vector = Tensor([10, 20]) # Shape: (2,)
>>> result = matrix + vector # Broadcasting: (2,2) + (2,) → (2,2)
>>> print(result.data)
[[11. 22.]
[13. 24.]]
HINTS:
- Use isinstance() to check if other is a Tensor
- NumPy handles broadcasting automatically with +
- Always return a new Tensor, don't modify self
- Preserve gradient tracking for future modules
"""
### BEGIN SOLUTION
if isinstance(other, Tensor):
# Tensor + Tensor: let NumPy handle broadcasting
return Tensor(self.data + other.data)
else:
# Tensor + scalar: NumPy broadcasts automatically
return Tensor(self.data + other)
### END SOLUTION
# nbgrader={"grade": false, "grade_id": "more-arithmetic", "solution": true}
def __sub__(self, other):
"""
Subtract two tensors element-wise.
Common use: Centering data (x - mean), computing differences for loss functions.
"""
### BEGIN SOLUTION
if isinstance(other, Tensor):
return Tensor(self.data - other.data)
else:
return Tensor(self.data - other)
### END SOLUTION
def __mul__(self, other):
"""
Multiply two tensors element-wise (NOT matrix multiplication).
Common use: Scaling features, applying masks, gating mechanisms in neural networks.
Note: This is * operator, not @ (which will be matrix multiplication).
"""
### BEGIN SOLUTION
if isinstance(other, Tensor):
return Tensor(self.data * other.data)
else:
return Tensor(self.data * other)
### END SOLUTION
def __truediv__(self, other):
"""
Divide two tensors element-wise.
Common use: Normalization (x / std), converting counts to probabilities.
"""
### BEGIN SOLUTION
if isinstance(other, Tensor):
return Tensor(self.data / other.data)
else:
return Tensor(self.data / other)
### END SOLUTION
# nbgrader={"grade": false, "grade_id": "matmul-impl", "solution": true}
def matmul(self, other):
"""
Matrix multiplication of two tensors.
TODO: Implement matrix multiplication using np.dot with proper validation
APPROACH:
1. Validate inputs are Tensors
2. Check dimension compatibility (inner dimensions must match)
3. Use np.dot for optimized computation
4. Return new Tensor with result
EXAMPLE:
>>> a = Tensor([[1, 2], [3, 4]]) # 2×2
>>> b = Tensor([[5, 6], [7, 8]]) # 2×2
>>> result = a.matmul(b) # 2×2 result
>>> # Result: [[1×5+2×7, 1×6+2×8], [3×5+4×7, 3×6+4×8]] = [[19, 22], [43, 50]]
SHAPE RULES:
- (M, K) @ (K, N) → (M, N) ✓ Valid
- (M, K) @ (J, N) → Error ✗ K ≠ J
COMPLEXITY: O(M×N×K) for (M×K) @ (K×N) matrices
HINTS:
- np.dot handles the optimization for us
- Check self.shape[-1] == other.shape[-2] for compatibility
- Provide clear error messages for debugging
"""
### BEGIN SOLUTION
if not isinstance(other, Tensor):
raise TypeError(f"Expected Tensor for matrix multiplication, got {type(other)}")
# Handle edge cases
if self.shape == () or other.shape == ():
# Scalar multiplication
return Tensor(self.data * other.data)
# For matrix multiplication, we need at least 1D tensors
if len(self.shape) == 0 or len(other.shape) == 0:
return Tensor(self.data * other.data)
# Check dimension compatibility for matrix multiplication
if len(self.shape) >= 2 and len(other.shape) >= 2:
if self.shape[-1] != other.shape[-2]:
raise ValueError(
f"Cannot perform matrix multiplication: {self.shape} @ {other.shape}. "
f"Inner dimensions must match: {self.shape[-1]} ≠ {other.shape[-2]}. "
f"💡 HINT: For (M,K) @ (K,N) → (M,N), the K dimensions must be equal."
)
elif len(self.shape) == 1 and len(other.shape) == 2:
# Vector @ Matrix
if self.shape[0] != other.shape[0]:
raise ValueError(
f"Cannot multiply vector {self.shape} with matrix {other.shape}. "
f"Vector length {self.shape[0]} must match matrix rows {other.shape[0]}."
)
elif len(self.shape) == 2 and len(other.shape) == 1:
# Matrix @ Vector
if self.shape[1] != other.shape[0]:
raise ValueError(
f"Cannot multiply matrix {self.shape} with vector {other.shape}. "
f"Matrix columns {self.shape[1]} must match vector length {other.shape[0]}."
)
# Perform optimized matrix multiplication
# Use np.matmul (not np.dot) for proper batched matrix multiplication with 3D+ tensors
result_data = np.matmul(self.data, other.data)
return Tensor(result_data)
### END SOLUTION
# nbgrader={"grade": false, "grade_id": "shape-ops", "solution": true}
def reshape(self, *shape):
"""
Reshape tensor to new dimensions.
TODO: Implement tensor reshaping with validation
APPROACH:
1. Handle different calling conventions: reshape(2, 3) vs reshape((2, 3))
2. Validate total elements remain the same
3. Use NumPy's reshape for the actual operation
4. Return new Tensor (keep immutability)
EXAMPLE:
>>> tensor = Tensor([1, 2, 3, 4, 5, 6]) # Shape: (6,)
>>> reshaped = tensor.reshape(2, 3) # Shape: (2, 3)
>>> print(reshaped.data)
[[1. 2. 3.]
[4. 5. 6.]]
COMMON USAGE:
>>> # Flatten for MLP input
>>> image = Tensor(np.random.rand(3, 32, 32)) # (channels, height, width)
>>> flattened = image.reshape(-1) # (3072,) - all pixels in vector
>>>
>>> # Prepare batch for convolution
>>> batch = Tensor(np.random.rand(32, 784)) # (batch, features)
>>> images = batch.reshape(32, 1, 28, 28) # (batch, channels, height, width)
HINTS:
- Handle both reshape(2, 3) and reshape((2, 3)) calling styles
- Check np.prod(new_shape) == self.size for validation
- Use descriptive error messages for debugging
"""
### BEGIN SOLUTION
# Handle both reshape(2, 3) and reshape((2, 3)) calling conventions
if len(shape) == 1 and isinstance(shape[0], (tuple, list)):
new_shape = tuple(shape[0])
else:
new_shape = shape
# Handle -1 for automatic dimension inference (like NumPy)
if -1 in new_shape:
if new_shape.count(-1) > 1:
raise ValueError("Can only specify one unknown dimension with -1")
# Calculate the unknown dimension
known_size = 1
unknown_idx = new_shape.index(-1)
for i, dim in enumerate(new_shape):
if i != unknown_idx:
known_size *= dim
unknown_dim = self.size // known_size
new_shape = list(new_shape)
new_shape[unknown_idx] = unknown_dim
new_shape = tuple(new_shape)
# Validate total elements remain the same
if np.prod(new_shape) != self.size:
raise ValueError(
f"Cannot reshape tensor of size {self.size} to shape {new_shape}. "
f"Total elements must match: {self.size} ≠ {np.prod(new_shape)}. "
f"💡 HINT: Make sure new_shape dimensions multiply to {self.size}"
)
# Reshape the data (NumPy handles the memory layout efficiently)
reshaped_data = np.reshape(self.data, new_shape)
# Preserve gradient tracking from the original tensor (important for autograd!)
result = Tensor(reshaped_data, requires_grad=self.requires_grad)
return result
### END SOLUTION
def transpose(self, dim0=None, dim1=None):
"""
Transpose tensor dimensions.
TODO: Implement tensor transposition
APPROACH:
1. Handle default case (transpose last two dimensions)
2. Handle specific dimension swapping
3. Use NumPy's transpose with proper axis specification
4. Return new Tensor
EXAMPLE:
>>> matrix = Tensor([[1, 2, 3], [4, 5, 6]]) # (2, 3)
>>> transposed = matrix.transpose() # (3, 2)
>>> print(transposed.data)
[[1. 4.]
[2. 5.]
[3. 6.]]
NEURAL NETWORK USAGE:
>>> # Weight matrix transpose for backward pass
>>> W = Tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]) # (3, 2)
>>> W_T = W.transpose() # (2, 3) - for gradient computation
>>>
>>> # Attention mechanism
>>> Q = Tensor([[1, 2], [3, 4]]) # queries (2, 2)
>>> K = Tensor([[5, 6], [7, 8]]) # keys (2, 2)
>>> attention_scores = Q.matmul(K.transpose()) # Q @ K^T
HINTS:
- Default: transpose last two dimensions (most common case)
- Use np.transpose() with axes parameter
- Handle 1D tensors gracefully (transpose is identity)
"""
### BEGIN SOLUTION
if dim0 is None and dim1 is None:
# Default: transpose last two dimensions
if len(self.shape) < 2:
# For 1D tensors, transpose is identity operation
return Tensor(self.data.copy())
else:
# Transpose last two dimensions (most common in ML)
axes = list(range(len(self.shape)))
axes[-2], axes[-1] = axes[-1], axes[-2]
transposed_data = np.transpose(self.data, axes)
else:
# Specific dimensions to transpose
if dim0 is None or dim1 is None:
raise ValueError("Both dim0 and dim1 must be specified for specific dimension transpose")
# Validate dimensions exist
if dim0 >= len(self.shape) or dim1 >= len(self.shape) or dim0 < 0 or dim1 < 0:
raise ValueError(
f"Dimension out of range for tensor with shape {self.shape}. "
f"Got dim0={dim0}, dim1={dim1}, but tensor has {len(self.shape)} dimensions."
)
# Create axes list and swap the specified dimensions
axes = list(range(len(self.shape)))
axes[dim0], axes[dim1] = axes[dim1], axes[dim0]
transposed_data = np.transpose(self.data, axes)
# Preserve requires_grad for gradient tracking (Module 05 will add _grad_fn)
result = Tensor(transposed_data, requires_grad=self.requires_grad if hasattr(self, 'requires_grad') else False)
return result
### END SOLUTION
# nbgrader={"grade": false, "grade_id": "reduction-ops", "solution": true}
def sum(self, axis=None, keepdims=False):
"""
Sum tensor along specified axis.
TODO: Implement tensor sum with axis control
APPROACH:
1. Use NumPy's sum with axis parameter
2. Handle axis=None (sum all elements) vs specific axis
3. Support keepdims to maintain shape for broadcasting
4. Return new Tensor with result
EXAMPLE:
>>> tensor = Tensor([[1, 2], [3, 4]])
>>> total = tensor.sum() # Sum all elements: 10
>>> col_sum = tensor.sum(axis=0) # Sum columns: [4, 6]
>>> row_sum = tensor.sum(axis=1) # Sum rows: [3, 7]
NEURAL NETWORK USAGE:
>>> # Batch loss computation
>>> batch_losses = Tensor([0.1, 0.3, 0.2, 0.4]) # Individual losses
>>> total_loss = batch_losses.sum() # Total: 1.0
>>> avg_loss = batch_losses.mean() # Average: 0.25
>>>
>>> # Global average pooling
>>> feature_maps = Tensor(np.random.rand(32, 256, 7, 7)) # (batch, channels, h, w)
>>> global_features = feature_maps.sum(axis=(2, 3)) # (batch, channels)
HINTS:
- np.sum handles all the complexity for us
- axis=None sums all elements (returns scalar)
- axis=0 sums along first dimension, axis=1 along second, etc.
- keepdims=True preserves dimensions for broadcasting
"""
### BEGIN SOLUTION
result = np.sum(self.data, axis=axis, keepdims=keepdims)
return Tensor(result)
### END SOLUTION
def mean(self, axis=None, keepdims=False):
"""
Compute mean of tensor along specified axis.
Common usage: Batch normalization, loss averaging, global pooling.
"""
### BEGIN SOLUTION
result = np.mean(self.data, axis=axis, keepdims=keepdims)
return Tensor(result)
### END SOLUTION
def max(self, axis=None, keepdims=False):
"""
Find maximum values along specified axis.
Common usage: Max pooling, finding best predictions, activation clipping.
"""
### BEGIN SOLUTION
result = np.max(self.data, axis=axis, keepdims=keepdims)
return Tensor(result)
### END SOLUTION
# nbgrader={"grade": false, "grade_id": "gradient-placeholder", "solution": true}
def backward(self):
"""
Compute gradients (implemented in Module 05: Autograd).
TODO: Placeholder implementation for gradient computation
STUDENT NOTE:
This method exists but does nothing until Module 05: Autograd.
Don't worry about it for now - focus on the basic tensor operations.
In Module 05, we'll implement:
- Gradient computation via chain rule
- Automatic differentiation
- Backpropagation through operations
- Computation graph construction
FUTURE IMPLEMENTATION PREVIEW:
```python
def backward(self, gradient=None):
# Module 05 will implement:
# 1. Set gradient for this tensor
# 2. Propagate to parent operations
# 3. Apply chain rule recursively
# 4. Accumulate gradients properly
pass
```
CURRENT BEHAVIOR:
>>> x = Tensor([1, 2, 3], requires_grad=True)
>>> y = x * 2
>>> y.sum().backward() # Calls this method - does nothing
>>> print(x.grad) # Still None
None
"""
### BEGIN SOLUTION
# Placeholder - will be implemented in Module 05
# For now, just ensure it doesn't crash when called
# This allows students to experiment with gradient syntax
# without getting confusing errors about missing methods
pass
### END SOLUTIONIn [ ]:
def test_unit_tensor_creation():
"""🧪 Test Tensor creation with various data types."""
print("🧪 Unit Test: Tensor Creation...")
# Test scalar creation
scalar = Tensor(5.0)
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
vector = Tensor([1, 2, 3])
assert np.array_equal(vector.data, np.array([1, 2, 3], dtype=np.float32))
assert vector.shape == (3,)
assert vector.size == 3
# Test matrix creation
matrix = Tensor([[1, 2], [3, 4]])
assert np.array_equal(matrix.data, np.array([[1, 2], [3, 4]], dtype=np.float32))
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
print("✅ Tensor creation works correctly!")
if __name__ == "__main__":
test_unit_tensor_creation()In [ ]:
def test_unit_arithmetic_operations():
"""🧪 Test arithmetic operations with broadcasting."""
print("🧪 Unit Test: Arithmetic Operations...")
# Test tensor + tensor
a = Tensor([1, 2, 3])
b = Tensor([4, 5, 6])
result = a + b
assert np.array_equal(result.data, np.array([5, 7, 9], dtype=np.float32))
# Test tensor + scalar (very common in ML)
result = a + 10
assert np.array_equal(result.data, np.array([11, 12, 13], dtype=np.float32))
# Test broadcasting with different shapes (matrix + vector)
matrix = Tensor([[1, 2], [3, 4]])
vector = Tensor([10, 20])
result = matrix + vector
expected = np.array([[11, 22], [13, 24]], dtype=np.float32)
assert np.array_equal(result.data, expected)
# Test subtraction (data centering)
result = b - a
assert np.array_equal(result.data, np.array([3, 3, 3], dtype=np.float32))
# Test multiplication (scaling)
result = a * 2
assert np.array_equal(result.data, np.array([2, 4, 6], dtype=np.float32))
# Test division (normalization)
result = b / 2
assert np.array_equal(result.data, np.array([2.0, 2.5, 3.0], dtype=np.float32))
# Test chaining operations (common in ML pipelines)
normalized = (a - 2) / 2 # Center and scale
expected = np.array([-0.5, 0.0, 0.5], dtype=np.float32)
assert np.allclose(normalized.data, expected)
print("✅ Arithmetic operations work correctly!")
if __name__ == "__main__":
test_unit_arithmetic_operations()In [ ]:
def test_unit_matrix_multiplication():
"""🧪 Test matrix multiplication operations."""
print("🧪 Unit Test: Matrix Multiplication...")
# Test 2×2 matrix multiplication (basic case)
a = Tensor([[1, 2], [3, 4]]) # 2×2
b = Tensor([[5, 6], [7, 8]]) # 2×2
result = a.matmul(b)
# Expected: [[1×5+2×7, 1×6+2×8], [3×5+4×7, 3×6+4×8]] = [[19, 22], [43, 50]]
expected = np.array([[19, 22], [43, 50]], dtype=np.float32)
assert np.array_equal(result.data, expected)
# Test rectangular matrices (common in neural networks)
c = Tensor([[1, 2, 3], [4, 5, 6]]) # 2×3 (like batch_size=2, features=3)
d = Tensor([[7, 8], [9, 10], [11, 12]]) # 3×2 (like features=3, outputs=2)
result = c.matmul(d)
# Expected: [[1×7+2×9+3×11, 1×8+2×10+3×12], [4×7+5×9+6×11, 4×8+5×10+6×12]]
expected = np.array([[58, 64], [139, 154]], dtype=np.float32)
assert np.array_equal(result.data, expected)
# Test matrix-vector multiplication (common in forward pass)
matrix = Tensor([[1, 2, 3], [4, 5, 6]]) # 2×3
vector = Tensor([1, 2, 3]) # 3×1 (conceptually)
result = matrix.matmul(vector)
# Expected: [1×1+2×2+3×3, 4×1+5×2+6×3] = [14, 32]
expected = np.array([14, 32], dtype=np.float32)
assert np.array_equal(result.data, expected)
# Test shape validation - should raise clear error
try:
incompatible_a = Tensor([[1, 2]]) # 1×2
incompatible_b = Tensor([[1], [2], [3]]) # 3×1
incompatible_a.matmul(incompatible_b) # 1×2 @ 3×1 should fail (2 ≠ 3)
assert False, "Should have raised ValueError for incompatible shapes"
except ValueError as e:
assert "Inner dimensions must match" in str(e)
assert "2 ≠ 3" in str(e) # Should show specific dimensions
print("✅ Matrix multiplication works correctly!")
if __name__ == "__main__":
test_unit_matrix_multiplication()In [ ]:
def test_unit_shape_manipulation():
"""🧪 Test reshape and transpose operations."""
print("🧪 Unit Test: Shape Manipulation...")
# Test basic reshape (flatten → matrix)
tensor = Tensor([1, 2, 3, 4, 5, 6]) # Shape: (6,)
reshaped = tensor.reshape(2, 3) # Shape: (2, 3)
assert reshaped.shape == (2, 3)
expected = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
assert np.array_equal(reshaped.data, expected)
# Test reshape with tuple (alternative calling style)
reshaped2 = tensor.reshape((3, 2)) # Shape: (3, 2)
assert reshaped2.shape == (3, 2)
expected2 = np.array([[1, 2], [3, 4], [5, 6]], dtype=np.float32)
assert np.array_equal(reshaped2.data, expected2)
# Test reshape with -1 (automatic dimension inference)
auto_reshaped = tensor.reshape(2, -1) # Should infer -1 as 3
assert auto_reshaped.shape == (2, 3)
# Test reshape validation - should raise error for incompatible sizes
try:
tensor.reshape(2, 2) # 6 elements can't fit in 2×2=4
assert False, "Should have raised ValueError"
except ValueError as e:
assert "Total elements must match" in str(e)
assert "6 ≠ 4" in str(e)
# Test matrix transpose (most common case)
matrix = Tensor([[1, 2, 3], [4, 5, 6]]) # (2, 3)
transposed = matrix.transpose() # (3, 2)
assert transposed.shape == (3, 2)
expected = np.array([[1, 4], [2, 5], [3, 6]], dtype=np.float32)
assert np.array_equal(transposed.data, expected)
# Test 1D transpose (should be identity)
vector = Tensor([1, 2, 3])
vector_t = vector.transpose()
assert np.array_equal(vector.data, vector_t.data)
# Test specific dimension transpose
tensor_3d = Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) # (2, 2, 2)
swapped = tensor_3d.transpose(0, 2) # Swap first and last dimensions
assert swapped.shape == (2, 2, 2) # Same shape but data rearranged
# Test neural network reshape pattern (flatten for MLP)
batch_images = Tensor(np.random.rand(2, 3, 4)) # (batch=2, height=3, width=4)
flattened = batch_images.reshape(2, -1) # (batch=2, features=12)
assert flattened.shape == (2, 12)
print("✅ Shape manipulation works correctly!")
if __name__ == "__main__":
test_unit_shape_manipulation()In [ ]:
def test_unit_reduction_operations():
"""🧪 Test reduction operations."""
print("🧪 Unit Test: Reduction Operations...")
matrix = Tensor([[1, 2, 3], [4, 5, 6]]) # Shape: (2, 3)
# Test sum all elements (common for loss computation)
total = matrix.sum()
assert total.data == 21.0 # 1+2+3+4+5+6
assert total.shape == () # Scalar result
# Test sum along axis 0 (columns) - batch dimension reduction
col_sum = matrix.sum(axis=0)
expected_col = np.array([5, 7, 9], dtype=np.float32) # [1+4, 2+5, 3+6]
assert np.array_equal(col_sum.data, expected_col)
assert col_sum.shape == (3,)
# Test sum along axis 1 (rows) - feature dimension reduction
row_sum = matrix.sum(axis=1)
expected_row = np.array([6, 15], dtype=np.float32) # [1+2+3, 4+5+6]
assert np.array_equal(row_sum.data, expected_row)
assert row_sum.shape == (2,)
# Test mean (average loss computation)
avg = matrix.mean()
assert np.isclose(avg.data, 3.5) # 21/6
assert avg.shape == ()
# Test mean along axis (batch normalization pattern)
col_mean = matrix.mean(axis=0)
expected_mean = np.array([2.5, 3.5, 4.5], dtype=np.float32) # [5/2, 7/2, 9/2]
assert np.allclose(col_mean.data, expected_mean)
# Test max (finding best predictions)
maximum = matrix.max()
assert maximum.data == 6.0
assert maximum.shape == ()
# Test max along axis (argmax-like operation)
row_max = matrix.max(axis=1)
expected_max = np.array([3, 6], dtype=np.float32) # [max(1,2,3), max(4,5,6)]
assert np.array_equal(row_max.data, expected_max)
# Test keepdims (important for broadcasting)
sum_keepdims = matrix.sum(axis=1, keepdims=True)
assert sum_keepdims.shape == (2, 1) # Maintains 2D shape
expected_keepdims = np.array([[6], [15]], dtype=np.float32)
assert np.array_equal(sum_keepdims.data, expected_keepdims)
# Test 3D reduction (simulating global average pooling)
tensor_3d = Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) # (2, 2, 2)
spatial_mean = tensor_3d.mean(axis=(1, 2)) # Average across spatial dimensions
assert spatial_mean.shape == (2,) # One value per batch item
print("✅ Reduction operations work correctly!")
if __name__ == "__main__":
test_unit_reduction_operations()In [ ]:
def test_module():
"""
Comprehensive test of entire module functionality.
This final test runs before module summary to ensure:
- All unit tests pass
- Functions work together correctly
- Module is ready for integration with TinyTorch
"""
print("🧪 RUNNING MODULE INTEGRATION TEST")
print("=" * 50)
# Run all unit tests
print("Running unit tests...")
test_unit_tensor_creation()
test_unit_arithmetic_operations()
test_unit_matrix_multiplication()
test_unit_shape_manipulation()
test_unit_reduction_operations()
print("\nRunning integration scenarios...")
# Test realistic neural network computation
print("🧪 Integration Test: Two-Layer Neural Network...")
# Create input data (2 samples, 3 features)
x = Tensor([[1, 2, 3], [4, 5, 6]])
# First layer: 3 inputs → 4 hidden units
W1 = Tensor([[0.1, 0.2, 0.3, 0.4],
[0.5, 0.6, 0.7, 0.8],
[0.9, 1.0, 1.1, 1.2]])
b1 = Tensor([0.1, 0.2, 0.3, 0.4])
# Forward pass: hidden = xW1 + b1
hidden = x.matmul(W1) + b1
assert hidden.shape == (2, 4), f"Expected (2, 4), got {hidden.shape}"
# Second layer: 4 hidden → 2 outputs
W2 = Tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6], [0.7, 0.8]])
b2 = Tensor([0.1, 0.2])
# Output layer: output = hiddenW2 + b2
output = hidden.matmul(W2) + b2
assert output.shape == (2, 2), f"Expected (2, 2), got {output.shape}"
# Verify data flows correctly (no NaN, reasonable values)
assert not np.isnan(output.data).any(), "Output contains NaN values"
assert np.isfinite(output.data).all(), "Output contains infinite values"
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])
# Reshape to 3D tensor (simulating batch processing)
tensor_3d = data.reshape(2, 2, 3) # (batch=2, height=2, width=3)
assert tensor_3d.shape == (2, 2, 3)
# Global average pooling simulation
pooled = tensor_3d.mean(axis=(1, 2)) # Average across spatial dimensions
assert pooled.shape == (2,), f"Expected (2,), got {pooled.shape}"
# Flatten for MLP
flattened = tensor_3d.reshape(2, -1) # (batch, features)
assert flattened.shape == (2, 6)
# Transpose for different operations
transposed = tensor_3d.transpose() # Should transpose last two dims
assert transposed.shape == (2, 3, 2)
print("✅ Complex shape operations work!")
# Test broadcasting edge cases
print("🧪 Integration Test: Broadcasting Edge Cases...")
# Scalar broadcasting
scalar = Tensor(5.0)
vector = Tensor([1, 2, 3])
result = scalar + vector # Should broadcast scalar to vector shape
expected = np.array([6, 7, 8], dtype=np.float32)
assert np.array_equal(result.data, expected)
# Matrix + vector broadcasting
matrix = Tensor([[1, 2], [3, 4]])
vec = Tensor([10, 20])
result = matrix + vec
expected = np.array([[11, 22], [13, 24]], dtype=np.float32)
assert np.array_equal(result.data, expected)
print("✅ Broadcasting edge cases work!")
print("\n" + "=" * 50)
print("🎉 ALL TESTS PASSED! Module ready for export.")
print("Run: tito module complete 01_tensor")
# Run comprehensive module test
if __name__ == "__main__":
test_module()In [ ]:
# YOUR ANSWER:
#
# Tensor A (1000, 1000) in float32:
# - Elements: ___________
# - Memory: ___________ MB
#
# Tensor B (500, 2000) in float32:
# - Elements: ___________
# - Memory: ___________ MB
#
# Which uses more memory? ___________
# How much more? ___________ MB
### BEGIN SOLUTION
# Tensor A: 1000 × 1000 = 1,000,000 elements
# Memory: 1,000,000 × 4 bytes = 4,000,000 bytes = 3.81 MB
# Tensor B: 500 × 2000 = 1,000,000 elements
# Memory: 1,000,000 × 4 bytes = 4,000,000 bytes = 3.81 MB
# Answer: Same memory usage (both have 1M elements)
# Difference: 0 MB - shape doesn't matter, only total elements
### END SOLUTIONIn [ ]:
# YOUR ANSWER:
#
# Layer 1 FLOPs: ___________
# Layer 2 FLOPs: ___________
# Layer 3 FLOPs: ___________
# Total FLOPs: ___________ (in millions)
### BEGIN SOLUTION
# Layer 1: (64, 784) @ (784, 256)
# FLOPs = 2 × 64 × 784 × 256 = 25,690,112
# Layer 2: (64, 256) @ (256, 128)
# FLOPs = 2 × 64 × 256 × 128 = 4,194,304
# Layer 3: (64, 128) @ (128, 10)
# FLOPs = 2 × 64 × 128 × 10 = 163,840
# Total: 25,690,112 + 4,194,304 + 163,840 = 30,048,256 FLOPs
# ≈ 30 million FLOPs per forward pass
### END SOLUTIONIn [ ]:
# YOUR ANSWER:
#
# A + B → Shape: ___________
# A + C → Shape: ___________
# A + D → Shape: ___________
# B + C → Shape: ___________
# C + D → Shape: ___________
### BEGIN SOLUTION
# A + B: (32, 64) + (64,) → (32, 64) [broadcast B to each row]
# A + C: (32, 64) + (32, 1) → (32, 64) [broadcast C to each column]
# A + D: (32, 64) + (1, 64) → (32, 64) [broadcast D to each row]
# B + C: (64,) + (32, 1) → (32, 64) [both broadcast to 2D]
# C + D: (32, 1) + (1, 64) → (32, 64) [outer product-like broadcast]
### END SOLUTIONIn [ ]:
# YOUR ANSWER:
#
# 1. FLOPs scaling when batch doubles: ___________
# (same / 2× / 4× / 8×?)
#
# 2. Memory scaling with float16 vs float32: ___________
# (same / 0.5× / 0.25× / 2×?)
#
# 3. Performance bottleneck: ___________
# (computation / memory bandwidth / both?)
#
# Reasoning: ___________
### BEGIN SOLUTION
# 1. FLOPs scale linearly with batch size: 2× FLOPs
# Original: 2 × 32 × 512 × 1024 = 33,554,432 FLOPs
# Doubled: 2 × 64 × 512 × 1024 = 67,108,864 FLOPs (2× increase)
# 2. Memory scales with precision: 0.5× memory (half the bytes per element)
# float32: 4 bytes/element
# float16: 2 bytes/element (50% reduction)
# 3. Bottleneck: Memory bandwidth for large batch sizes
# - Modern GPUs have high FLOP/s (teraFLOPs)
# - Memory bandwidth is limited (100s of GB/s)
# - Large matrices → more data movement than computation
# - For small batches: computation bound
# - For large batches: memory bandwidth bound
### END SOLUTION