Test a PyTorch-to-JAX migration by comparing both forward values and parameter gradients against an independently calculated answer. A CPU-only check can reject a broken derivative before you spend time on target-hardware testing, even when the migrated loss looks correct.
For example, take a scalar weight w=1, inputs x=[1, 2], targets y=[0, 0], and mean squared error. Both implementations should return a loss of 2.5 and a gradient of 5. Applying JAX's stop_gradient to the weight preserves that loss but changes its gradient to zero.
Matching losses can hide a broken update direction.
Use the local check to reject defects, not certify the migration
The XStack abstract describes a repository-migration benchmark in which agents receive a PyTorch repository and a clean TPU machine and must deliver a JAX repository. According to the authors, acceptance depends on grading criteria based on the source repository's published result, authenticated and calibrated by them.
That makes a concrete engineering decision important: which defects should stop a migration before repository-level evaluation begins?
The check below is an independent proposal, not a description of XStack's testing implementation. Its purpose is narrow: verify one differentiable computation in isolation. It neither reproduces the benchmark's grading nor establishes that a repository trains correctly on a TPU.
I would make this a required early gate for a migrated differentiable component, whether the translation came from an agent or a human. I would not use it as the migration's acceptance test.
Start with a derivative you can calculate independently
Use the objective:
L(w) = mean((w*x - y)^2)
With the example inputs:
L(1) = (1² + 2²) / 2 = 2.5
dL/dw = mean(2*x*(w*x - y))
= (2 + 8) / 2
= 5
These analytical values are the oracle. Comparing frameworks alone is weaker: two implementations can agree because they share the same mistake.
The migration defect is a single line:
w = jax.lax.stop_gradient(w)
During forward evaluation, that line leaves the weight's value unchanged. During differentiation, it blocks the path from the loss back to the input weight. The resulting loss is constant with respect to that input as far as automatic differentiation is concerned, so its gradient is zero.
A forward assertion has no reason to fail here. The derivative assertion does.
Build the isolated CPU test
The following recipe targets Python 3.11 on Linux x86-64. It pins the top-level packages and JAX's backend package; it is not a complete transitive dependency lock. The example has not been executed here, so treat these as reproduction instructions rather than a reported test run.
Create a temporary workspace and virtual environment:
workdir="$(mktemp -d)"
cd "$workdir"
python3.11 -m venv .venv
. .venv/bin/activate
python -m pip install \
--index-url https://download.pytorch.org/whl/cpu \
'torch==2.5.1+cpu'
python -m pip install \
'numpy==1.26.4' \
'jax==0.4.35' \
'jaxlib==0.4.35' \
'pytest==8.3.3'
Save this as test_gradient_parity.py in that directory:
import jax
import jax.numpy as jnp
import numpy as np
import torch
RTOL = 1e-5
ATOL = 1e-6
def torch_loss(w, x, y):
return ((w * x - y) ** 2).mean()
def jax_loss(w, x, y):
# Deliberate migration defect: forward value survives,
# but the derivative with respect to w is blocked.
w = jax.lax.stop_gradient(w)
return jnp.mean((w * x - y) ** 2)
def assert_close(actual, expected, label):
np.testing.assert_allclose(
actual,
expected,
rtol=RTOL,
atol=ATOL,
err_msg=label,
)
def test_loss_and_gradient_parity():
tx = torch.tensor([1.0, 2.0], dtype=torch.float32)
ty = torch.tensor([0.0, 0.0], dtype=torch.float32)
tw = torch.tensor(1.0, dtype=torch.float32, requires_grad=True)
t_loss = torch_loss(tw, tx, ty)
(t_grad,) = torch.autograd.grad(t_loss, tw)
jx = jnp.array([1.0, 2.0], dtype=jnp.float32)
jy = jnp.array([0.0, 0.0], dtype=jnp.float32)
jw = jnp.array(1.0, dtype=jnp.float32)
j_loss, j_grad = jax.value_and_grad(jax_loss)(jw, jx, jy)
# Convert to host scalars before making assertions.
tl = t_loss.detach().item()
tg = t_grad.detach().item()
jl = float(j_loss)
jg = float(j_grad)
assert np.isfinite([tl, tg, jl, jg]).all(), "non-finite result"
# Independent analytical checks, not just framework agreement.
assert_close(tl, 2.5, "PyTorch loss oracle")
assert_close(jl, 2.5, "JAX loss oracle")
assert_close(jl, tl, "forward parity")
assert_close(tg, 5.0, "PyTorch gradient oracle")
assert_close(jg, 5.0, "JAX gradient oracle")
assert_close(jg, tg, "gradient parity")
Run it with CPU execution and JAX's 64-bit mode disabled explicitly:
CUDA_VISIBLE_DEVICES="" \
JAX_PLATFORMS=cpu \
JAX_ENABLE_X64=false \
python -m pytest -q test_gradient_parity.py
The expected failure is the assertion labeled JAX gradient oracle: the defective function yields a gradient of 0, while the oracle requires 5. The preceding forward checks should pass. This expected result follows from the calculation above, not from an observed execution log.
The explicit tolerances allow small numerical differences without accepting this defect. They are starting values for this float32 scalar example, not universal tolerances for a migrated model.
Repair the derivative, then keep the failing case
Remove the gradient barrier:
def jax_loss(w, x, y):
return jnp.mean((w * x - y) ** 2)
Rerun the same command. The repaired expression has the analytical loss 2.5 and derivative 5, so all assertions are expected to pass.
Keep the original failure as a regression fixture for the check itself. Before relying on a verification gate, temporarily reintroduce the defect and confirm that the gate rejects it. Then restore the repair.
A gate should demonstrate that it can reject its intended defect.
For this deterministic fixture, there is no reason to introduce random initialization, a data loader, or an optimizer. Those would add state without improving coverage of the specific gradient barrier.
When moving beyond the scalar example, compare gradients by named parameter and assert matching shapes before comparing values. Do not flatten differently structured parameter trees and assume that positional agreement means semantic agreement.
Keep three verification layers separate
Use this checklist to decide what a passing check actually permits:
| Layer | What to check | What a pass permits |
|---|---|---|
| Local semantic parity | Fixed inputs, explicit dtypes, expected losses, gradient shapes and values, finite results | Proceed past the covered computation |
| Representative training behavior | Controlled initial parameters and batches, optimizer-state mapping, update deltas, short loss trajectories | Proceed past the covered training path |
| Target-device validation | The intended JAX transformations, device placement, precision settings, distributed execution, and repository acceptance criteria | Accept only the tested target configuration and criteria |
CPU parity is an early filter, not a hardware verdict.
For a real component, choose fixtures that exercise its meaningful branches. If the migration changes a reduction axis, include an asymmetric input that distinguishes the axes. If it changes masking, include masked and unmasked elements. An all-zero fixture can make a broken derivative look correct simply because the expected derivative is also zero.
Likewise, test the execution path you intend to deploy. An eager value_and_grad check does not exercise a separate JIT-compiled or distributed training wrapper. Add checks for those paths rather than assuming the scalar result covers them.
Static analysis can supplement this work by flagging suspicious gradient barriers for review. A barrier's presence alone is not a defect: some algorithms intentionally block gradients. The numerical contract determines whether it is wrong here.
Try one known-answer fixture today
In five minutes, choose one migrated scalar loss or small tensor operation, write down its expected derivative, and assert both frameworks against that answer. Deliberately block one required gradient path and verify that your check fails.
That exercise establishes coverage of one computation, not repository equivalence, training convergence, or TPU correctness. Its value is a small, explainable failure that stops a known defect before a more expensive evaluation.
Use that fixture as the starting point for putting a verification gate on your pipeline.