JAX PathwaysCheck your starting point
A classifier predicts the majority class on every example and reports high accuracy. What information is missing?
Compare your reasoning
Class counts and per-class errors are missing. A majority baseline can score highly on imbalanced data while missing every rare example. Inspect the confusion matrix and the denominator for each class.
If this is unfamiliar, work through Train a CNN on a small dataset before attempting the project.
Build in stages
1. Connect equations to module state
Check a forward pass and unreduced loss independently, then identify parameter and nonparameter state in the NNX model.
2. Train and evaluate under separate contracts
Keep held-out examples out of updates, verify evaluation leaves state unchanged and read confusion rows as true-class denominators.
3. Diagnose one failure at a time
Use recorded input ranges, losses, gradient norms and update sizes to distinguish a scale problem from a data or shape problem.
Try a changed condition
Make a prediction before opening the reasoning. Keep your answer and the evidence you would collect.
Training loss falls, but recall for a rare class is zero. Write the next three checks before changing the model.
Compare an approach
Inspect class counts and labels; compare a constant predictor and per-class confusion; then inspect score distributions and choose any threshold on validation data. Report final metrics on held-out data without using it to tune the repair.
Diagnose a likely failure
Symptom: Evaluation results change even though you did not call an optimizer.
Check next: Compare state snapshots and train/eval flags; stochastic or running-statistic behavior may still be active.
Decide what is ready
Use mlp-classifier before the larger image harness. Retain independent forward/loss checks, two runs with declared seeds, state-pure evaluation and class-specific failure analysis.
Where to put your project code
Write your code in: projects/mlp-classifier/my_model.py (copy projects/mlp-classifier/starter/model.py first: cp projects/mlp-classifier/starter/model.py projects/mlp-classifier/my_model.py).
Verify your implementation: python3 projects/mlp-classifier/tests/check.py --stage all --impl projects/mlp-classifier/my_model.py
Run the staged project checks on your implementation in projects/mlp-classifier/my_model.py and record your transfer experiment results in your portfolio.
Open Reproducible JAX/Flax experiments workshop
What still needs further work
Natural-image dataset ingestion, richer augmentation studies and pretrained vision transfer remain larger follow-up labs; synthetic images are separate from real-photo generalization.