JAX PathwaysCheck your starting point
A JAX call returns an array handle almost immediately. Does a timer ending there necessarily measure completed computation?
Compare your reasoning
No. Synchronize the result inside the measured boundary. Report first-call setup separately from warmed execution and keep inputs, dtype and output checks fixed.
If this is unfamiliar, work through Benchmark asynchronous work correctly before attempting the project.
Build in stages
1. Choose a comparable measurement
Record warmup, synchronization and sample distributions. Change shapes and host reads separately so you can identify their effects.
2. Explain the evidence
Locate one annotated operation in a real trace, distinguish nested from additive time and state the memory level behind a roofline estimate.
3. Test a tradeoff
Compare rematerialized and ordinary gradients before timing. Distinguish saved residuals, compiler estimates and observed peak memory.
Try a changed condition
Make a prediction before opening the reasoning. Keep your answer and the evidence you would collect.
A kernel becomes twice as fast but the full step barely changes. What does the result suggest?
Compare an approach
First measure the kernel’s original share of the step. Even eliminating a small component has limited total benefit. Also inspect copies, synchronization and changed compilation; a local kernel timer cannot explain the complete pipeline alone.
Diagnose a likely failure
Symptom: A claimed speedup disappears in repeated runs.
Check next: Check equivalent work, compilation reuse, synchronization, sample spread and competing system activity; retain the raw samples.
Decide what is ready
Use the trace and timing stage of sharded-training. Keep before/after correctness, named measurement boundaries, trace locations and distributions with a justified accept/reject decision.
Where to put your project code
Write your code in: projects/performance-audit/starter/model.py (copy projects/performance-audit/starter/model.py first: cp projects/performance-audit/starter/model.py projects/performance-audit/starter/model.py).
Verify your implementation: python3 projects/performance-audit/tests/check.py --stage all --impl projects/performance-audit/starter/model.py
Run the staged project checks on your implementation in projects/performance-audit/starter/model.py and record your transfer experiment results in your portfolio.
Open Resume and profile a sharded training run
What still needs further work
CPU traces do not characterize TPU/GPU communication or memory. Target traces and sustained workload measurements remain required for accelerator claims.