JAX PathwaysCheck your starting point
For a squared residual, is the gradient of the average loss the same shape as the stack of per-example gradients?
Compare your reasoning
No. A stack retains the examples axis; the gradient of the average reduces that axis. Their mean agrees for the same equally weighted differentiable objective.
If this is unfamiliar, work through Losses and value_and_grad before attempting the project.
Build in stages
1. Establish the derivative
Predict derivatives at several inputs and compare with an analytic calculation. Check that the loss is a scalar before differentiating.
2. Compose without changing the question
Compare a Python loop, mapped gradients and their compiled version on the same inputs. Change numerical values separately from shapes and static configuration; record tracing and timing as different observations.
Try a changed condition
Make a prediction before opening the reasoning. Keep your answer and the evidence you would collect.
One minibatch has two examples and another has six. How should you combine their mean gradients?
Compare an approach
Weight each mean gradient by its example count, add the weighted gradients, then divide by eight. An unweighted average gives examples in the smaller batch too much influence. Verify against a gradient over the concatenated batch.
Diagnose a likely failure
Symptom: A Python conditional works eagerly but raises a tracer error under jit.
Check next: Identify the data-dependent truth test. Use a JAX control-flow operation with compatible branch outputs; do not mark changing array data static to suppress the symptom.
Decide what is ready
Use foundation-toolkit stage 3. Keep independent derivatives, batched/compiled parity and one changed signature. Explain why a faster warm call does not measure first-call cost.
Where to put your project code
Write your code in: projects/transforms-audit/starter/model.py (copy projects/transforms-audit/starter/model.py first: cp projects/transforms-audit/starter/model.py projects/transforms-audit/starter/model.py).
Verify your implementation: python3 projects/transforms-audit/tests/check.py --stage all --impl projects/transforms-audit/starter/model.py
Run the staged project checks on your implementation in projects/transforms-audit/starter/model.py and record your transfer experiment results in your portfolio.
Open Connect arrays, transformations and explicit state
What still needs further work
Transformation combinations involving nondifferentiable boundaries require further examples; passing smooth fixtures does not validate arbitrary gradients.