PART 1
0/3PRNG & Random Sampling
Set up JAX's functional PRNG: wrap seeds, split keys, and sample random tensors.
MLP in JAX from Scratch
MLP in JAX from Scratch
Implement initialization, forward passes, loss, autodiff, and pure functional SGD updates in JAX.
Every source step is its own lesson with intuition, concepts, correctly rendered MathJax mathematics, implementation, tests, mistakes, and a checkpoint.
PART 2
0/3Synthetic Data & Labels
Generate input features, assign deterministic class labels, and one-hot encode targets.
PART 3
0/2Parameter Initialization
Initialize a single linear layer and extend to a full MLP parameter list.
PART 4
0/4Forward Pass Building Blocks
Implement the linear layer, ReLU, softmax, and the full MLP forward pass.
PART 5
0/3Loss & Metrics
Compute log-softmax, mean cross-entropy loss, and classification accuracy.
PART 6
0/3Autodiff & SGD Update
Wrap the loss for autodiff, take gradients with jax.grad, and apply SGD updates.
PART 7
0/3Training Loop & Inference
Combine the pieces into a training step, an epoch loop, and a prediction function.