T2J-Bench: Can Coding Agents Preserve Training and Inference Behavior in PyTorch-to-JAX Migration?

Background: PyTorch-to-JAX Migration Challenges

Migrating a PyTorch model repository to JAX requires coordinated rewriting of model definitions, weight loading, training loops, and inference caches. Small changes can alter gradient directions or training trajectories. The article illustrates how a simple reordering of operations (e.g., mean of squares vs. square of mean) changes loss and gradients, and similar issues appear in optimizer updates, attention masks, and KV caches.

T2J-Bench: Benchmark Design

T2J-Bench comprises 63 self-contained repository-level migration tasks covering 18 model families. Each task provides a PyTorch source tree (15–20 modules) and public weights. Agents must deliver a native JAX/Flax repository preserving file structure, public interfaces, and implementing PyTorch-free checkpoint conversion.

Training tasks: 45 LoRA fine-tuning tasks across 15 model families with SFT, DPO, PPO objectives.

Serving tasks: 18 tasks across 6 model families with single-sequence, variable-length batching, and chunked prefill paged-attention modes.

Evaluation uses three dimensions: Spec (structure, interfaces, framework independence), Numeric (module outputs, loss, gradients), and Behavioral (training dynamics, inference behavior). Candidate code runs in isolated, offline, PyTorch-free environments against pre-recorded PyTorch references. The final reward is a continuous 0–1 score with Spec as a multiplicative gate and Numeric/Behavioral equally weighted. A reference JAX implementation scores 0.980, confirming task solvability.

TwinHarness: Execution-Grounded Feedback Loop

TwinHarness leverages the runnable original repository as an "executable specification." It organizes agent work into three phases:

Diagnose: Read current code and existing run evidence to locate migration gaps and set repair targets.

Migrate: Modify JAX/Flax implementation, weight conversion, and training/serving paths.

Verify: Execute updated repository, record check results, and pass evidence to next diagnosis round. Code state and diagnostic evidence are persisted to avoid reliance on conversation memory.

Verification includes two signal types: (1) differential checks comparing PyTorch and JAX outputs under matched inputs, and (2) candidate-side checks (imports, parameter trees, gradient validity, JAX tracing, training entry execution). Each conclusion is labeled "proven," "violated," or "unproven." TwinHarness uses only task-visible resources, keeping development feedback separate from final evaluation.

Experimental Results

Five agent harnesses with supported models yielded 20 off-the-shelf configurations, each run on all 63 tasks with three random seeds. The strongest off-the-shelf configuration, OpenCode + opus-5, achieved an overall score of 0.676 (Spec 0.927, Numeric 0.670, Behavioral 0.741), revealing a gap between structural compliance and behavioral fidelity.

With Claude-Code + opus-5 fixed, TwinHarness raised the overall score from 0.649 to 0.723 (absolute gain 0.074, relative 11.4%), outperforming the best off-the-shelf baseline by ~7%. Numeric improved from 0.618 to 0.730; Behavioral from 0.728 to 0.792. Gains were not universal: OpenCode + sonnet-5 still led on serving tasks, and all configurations remained far from the 0.980 reference.

Quality–Cost Trade-offs

API cost per 63-task run was measured. At comparable cost, TwinHarness (0.723) outperformed additional vanilla continuation rounds of Claude-Code (~0.66) and OpenCode (~0.60). Initial OpenCode (0.676) offered lower cost; TwinHarness offered higher fidelity—both on the observed Pareto frontier. Notably, extra vanilla rounds degraded OpenCode (0.676 → 0.628 → 0.60), suggesting unstructured iterations can break correct behavior.

Generalization Across Harness–Backbone Pairs

Four matched harness–backbone pairs were tested, each showing improvement with TwinHarness:

Claude-Code + opus-5: 0.649 → 0.723

Claude-Code + sonnet-5: 0.408 → 0.467

Gemini-CLI + gemini-3.7-flash: 0.102 → 0.222

Gemini-CLI + gemini-3.8-flash: 0.087 → 0.217

Absolute gains ranged from 0.059 to 0.130, indicating broad applicability. The authors suggest integrating systematic source-execution feedback into existing agent pipelines.

Effect on Subsequent Repair Rounds

Comparing first-to-second round changes: vanilla OpenCode + opus-5 dropped from 0.676 to 0.628 (22 tasks regressed, 7 improved), while TwinHarness on OpenCode rose from 0.632 to 0.706 (11 improved, 2 regressed). This supports the hypothesis that structured verification helps preserve correct behavior, though component-level causality is not isolated and continued rounds may not sustain gains.

Limitations and Future Work

Experiments are limited to T2J-Bench matched configurations; generalization to other frameworks, repository types, and deployment conditions remains untested. The open-sourced benchmark ( https://github.com/google-research/t2j-bench/) enables unified comparison of agent systems and further investigation of which execution evidence best preserves cross-framework behavior.