JAX PathwaysCheck your starting point
Weights match after conversion, but predictions differ. What besides weights must be compared?
Compare your reasoning
Compare input normalization, feature/channel order, operation semantics, non-trainable state, output labels, dtype and shape. Start with a hand-computed asymmetric example to expose layout mistakes.
If this is unfamiliar, work through Move models between JAX, Keras, TensorFlow and PyTorch before attempting the project.
Build in stages
1. Choose the model and preserve its meaning
Adapt the model under a fixed held-out protocol before freezing its release. Compare framework layouts, preprocessing and state with independent predictions. Convert the actual PyTorch model to Flax and locate an injected operation mismatch numerically. Keep a candidate decision table, feature/parameter mapping table, and the first divergent operation from a deliberately broken conversion.
2. Freeze the artifact and numeric policy
Export and load in a fresh process, reject malformed inputs, then track weight, activation, accumulator, bias and output formats. Separate clipping from rounding and retain held-out conversion checks. Include a fresh-interpreter export receipt, raw-request rejection cases, clipping counts and actual versus ideal packed byte counts.
3. Deliver to a stated runtime
Package and qualify the selected service, relate queue assumptions to measured service work, then examine the additional conversion and device budgets for an edge target. Leave actual device fields unmeasured until tested. Keep separate artifact, process, image and device receipts. Explain an interpolated percentile and a burst deadline failure before choosing a capacity intervention.
Try a changed condition
Make a prediction before opening the reasoning. Keep your answer and the evidence you would collect.
An INT8 model matches typical inputs but fails after a sensor’s range changes. How do you locate the regression?
Compare an approach
Verify raw preprocessing first; count values outside the frozen calibration range, inspect clipped channels and compare float-converted versus quantized outputs. Do not recalibrate on final test labels. Re-evaluate a revised policy on a separate held-out set.
Diagnose a likely failure
Symptom: A smaller model file produces no lower request latency.
Check next: Inspect load-time expansion, fallback operators, transfers and preprocessing on the named runtime. File size is not a kernel or latency measurement.
Decide what is ready
Use deployment-audit and engineering-release to assemble one release dossier: candidate selection and retention criteria; a named feature and parameter mapping; source/export and layerwise parity; full precision policy and calibration identity; raw-request validation; artifact and process receipts; measured request samples; and a queue/deadline analysis. Mark container and target-device evidence unmeasured unless you actually ran those environments. Another learner should be able to reproduce each claimed check from the recorded inputs and command.
Where to put your project code
Write your code in: projects/deployment-audit/my_model.py (copy projects/deployment-audit/starter/model.py first: cp projects/deployment-audit/starter/model.py projects/deployment-audit/my_model.py).
Verify your implementation: python3 projects/deployment-audit/tests/check.py --stage all --impl projects/deployment-audit/my_model.py
Run the staged project checks on your implementation in projects/deployment-audit/my_model.py and record your transfer experiment results in your portfolio.
Open Ship a tracked and containerized model release
Open Convert PyTorch weights and audit Flax architecture parity
What still needs further work
Actual edge-device delegates, thermal/energy budgets, packed low-bit kernels and production serving loads remain unqualified. The revised sequence moves adaptation before export and containerization before service capacity planning.