# Let's trace through the shapes step by step
print("Shape flow through network:")
print(f"Input X:        {X.shape}")           # (4, 2)
print(f"Weights W1:     {W1.shape}")          # (2, 4)
print(f"X @ W1:         {(X @ W1).shape}")    # (4, 4)
print(f"Bias b1:        {b1.shape}")          # (1, 4)

# Broadcasting explanation
sample_mult = X @ W1
print(f"\nBroadcasting bias:")
print(f"(X @ W1) shape: {sample_mult.shape}")  # (4, 4)
print(f"b1 shape:       {b1.shape}")           # (1, 4)
print(f"Result shape:   {(sample_mult + b1).shape}")  # (4, 4)