What problem does it solve? Porting a reference ML or scientific computing implementation to JAX often produces code that looks equivalent but silently diverges in outputs, gradients, or post-update weights. This Skill enforces a parity-first methodology so the JAX port provably reproduces the reference's behavior instead of merely resembling it. ## Core Features & Use Cases - Replay-Based Parity Testing: Record one reference run (weights, batches, RNG decisions, intermediates, gradients) and drive the JAX side by replaying those decisions, isolating implementation differences from stochastic noise. - Bottom-Up Verification: Parity-test each layer or primitive before building the one above it, covering weights, modules, pure ops, loss, and the full train step. - Numerical-Equivalence Debugging: Distinguish real port bugs from float32 conditioning issues by recomputing disputed values in float64 and finite-differencing gradients. - Use Case: When porting a PyTorch training loop to JAX/Optax, record the reference run, convert weights via an explicit name-to-pytree map, and verify gradients and post-update weights match within a defined tolerance policy. ## Quick Start Port this PyTorch model to JAX and verify numerical parity against the recorded reference outputs, gradients, and updated weights.
