T2J-BENCH: Execution-Grounded PyTorch-to-JAX Repository Migration

The benchmark

Beyond code translation.
Preserving executable behavior.

Migrating PyTorch to JAX is a structural rewrite: stateful modules become pure, compilable functions with explicit parameters and state. T2J-BENCH tests whether that rewrite preserves the original system.

T2J-BENCH pipeline: curate 63 training and serving tasks, migrate PyTorch repositories and checkpoints to native JAX/Flax, then evaluate Spec, Numeric, and Behavioral equivalence in an isolated torch-free verifier. — Whole-repository migration, from model code and checkpoint conversion to end-to-end execution. Click the figure for the full-resolution PDF.

The central challenge

A repository can satisfy its interface and still diverge in forward numerics, gradients, optimizer trajectories, or decode behavior. T2J-BENCH measures those differences by running the code.

Task coverage

One migration contract.
Two sides of the ML lifecycle.

Every task starts with a self-contained PyTorch source repository and raw checkpoint. Agents deliver native JAX/Flax code, a torch-free checkpoint converter, and an executable training or serving path.

Training migration45 tasks

15 families × 3 methods

Preserve forward passes, losses, gradients, LoRA adapters, optimizer state, and the data pipeline.

SFTDPOPPO

24 text · 21 vision-language tasks

Explore the 15 model families

BLOOMFalconLlama 3MistralPhi-4Qwen2Qwen3StarCoder2Gemma 3InternVLLLaVA-NeXTPaliGemmaQwen2-VLQwen3-VLQwen2-Omni

Serving migration18 tasks

6 families × 3 engine modes

Preserve paged KV caches, prefill and decode behavior, sampling, and multi-request scheduling.

CoreRagged batchChunked prefill

15 text · 3 vision-language tasks

Explore the 6 model families

Llama 3Qwen2Qwen3GPT-OSSGemma-4Qwen2.5-VL

18 distinct families across both classes; three families appear in both. Evaluated modalities are text and vision, with no graded audio modality.

Execution-grounded evaluation

Run it. Compare it. Verify independence.

Each candidate runs in a separate, offline verifier with PyTorch removed. Ground truth is recorded from the original PyTorch repository; the reference JAX solution is evaluated as a candidate.

01 / Spec

Is it a genuine migration?

Check structure, public interfaces, and framework independence. Hidden PyTorch dependencies cannot earn equivalence credit.

02 / Numeric

Do the numbers agree?

Compare module outputs, composed forward passes, losses, gradients, and step-wise computations against frozen references.

03 / Behavioral

Does the system behave the same?

Evaluate training dynamics, held-out behavior, and paged decode execution under a uniform runtime budget.

R = s_spec × (0.5 s_num + 0.5 s_beh)

A spec-gated reward in [0, 1], with continuous, threshold-free numeric and behavioral similarity.

7 evaluation levels

L0–L6 cover independence and structure through end-to-end execution.

3 random seeds

Each configuration is evaluated across all 63 tasks; results are seed averages.

0.980 oracle reward

A known-good JAX solution establishes task solvability and remaining headroom.

Main results

High spec scores still leave
a substantial behavior gap.

The strongest of twenty off-the-shelf configurations reaches 0.676 reward. TwinHarness raises its matched Claude-Code / opus-5 baseline from 0.649 to 0.723, while the oracle reaches 0.980.

Claude-Code / opus-5

0.649

Best off-the-shelf

0.676

TwinHarness / opus-5
0.723

Reference JAX oracle

0.980

Three-seed mean scores on all 63 tasks. Reward shows sample standard deviation; higher is better. Train and Serve are subset rewards.
Harness	Model	Reward ↑	Spec	Numeric	Behavior	Train	Serve
TwinHarness Ours	opus-5	0.723 ± 0.022	0.931	0.730	0.792	0.764	0.622
OpenCode	opus-5	0.676 ± 0.014	0.927	0.670	0.741	0.747	0.499
Claude-Code	opus-5	0.649 ± 0.008	0.924	0.618	0.728	0.725	0.457
OpenHands	opus-5	0.622 ± 0.029	0.907	0.618	0.681	0.709	0.403
OpenCode	gemini-3.7-flash	0.526 ± 0.033	0.883	0.592	0.535	0.595	0.351
OpenCode	sonnet-5	0.472 ± 0.042	0.760	0.467	0.554	0.390	0.677
OpenCode	glm-5.2	0.470 ± 0.051	0.732	0.495	0.544	0.491	0.418
OpenHands	gemini-3.8-flash	0.443 ± 0.016	0.828	0.555	0.435	0.525	0.235
OpenHands	gemini-3.7-flash	0.416 ± 0.030	0.851	0.537	0.377	0.488	0.234
Claude-Code	sonnet-5	0.408 ± 0.032	0.732	0.424	0.479	0.329	0.605
OpenHands	sonnet-5	0.401 ± 0.023	0.755	0.428	0.459	0.321	0.602
Terminus-2	gemini-3.8-flash	0.369 ± 0.015	0.793	0.496	0.353	0.453	0.160
OpenCode	gemini-3.8-flash	0.361 ± 0.013	0.633	0.423	0.381	0.448	0.144
Terminus-2	gemini-3.7-flash	0.358 ± 0.012	0.786	0.479	0.331	0.442	0.147
OpenHands	glm-5.2	0.247 ± 0.026	0.572	0.279	0.279	0.223	0.306
Gemini-CLI	gemini-3.7-flash	0.102 ± 0.020	0.346	0.120	0.111	0.081	0.155
Gemini-CLI	gemini-3.8-flash	0.087 ± 0.012	0.324	0.100	0.102	0.069	0.131
OpenHands	gemini-3.1-pro	0.086 ± 0.031	0.364	0.108	0.092	0.040	0.203
OpenCode	gemini-3.1-pro	0.064 ± 0.023	0.355	0.069	0.083	0.034	0.139
Gemini-CLI	gemini-3.1-pro	0.057 ± 0.032	0.345	0.066	0.068	0.027	0.133
Terminus-2	gemini-3.1-pro	0.006 ± 0.004	0.223	0.008	0.008	0.001	0.019
Reference JAX oracle	—	0.980	1.000	0.998	0.961	0.973	0.997

+11.4%

Matched-baseline reward gain

TwinHarness improves Claude-Code / opus-5 by 0.074 absolute reward, from 0.649 to 0.723.

0.257

Remaining oracle gap

Even the best evaluated configuration leaves considerable room for more faithful migrations.

Uneven gains

Aggregate progress has limits

OpenCode / sonnet-5 leads serving reward at 0.677. TwinHarness does not lead every task subtype.

TwinHarness

Use the source as
an executable specification.

A reference-grounded diagnose–migrate–verify loop turns the task-visible PyTorch implementation into structured feedback for the next edit.

TwinHarness diagnose, migrate, and verify loop, comparing task-visible PyTorch behavior with the evolving JAX candidate and carrying diagnostic evidence across iterations. — Repository state and diagnostic evidence are maintained separately. Checks record what is established, violated, or not yet established.

A strict information boundary

TwinHarness uses the task specification, visible PyTorch source, released checkpoint, and candidate execution. It does not access hidden tests, frozen reference outputs, reference JAX solutions, verifier code, or benchmark rewards.

Structured refinement

Better feedback makes
the next iteration count.

With Claude-Code / opus-5, TwinHarness improves from 0.662 after one round to 0.723 after two. Ordinary continuation improves from 0.649 to 0.658.

With OpenCode, second-round regressions affect 2 of 63 tasks under TwinHarness versus 22 under ordinary continuation. Each method starts from its own first-round candidate.

One- versus two-round reward comparison for base agents and TwinHarness with Claude-Code and OpenCode, using opus-5. — Gains emerge through iteration; they are not uniform at every round budget.

Citation

BibTeX

@article{t2jbench2026,
  title={T2J-BENCH: Benchmarking Coding Agents on Execution-Grounded PyTorch-to-JAX Repository Migration},
  author={Feng, Tao and Chen, Jiefeng and Song, Linxin and Zhuang, Yufan and Huang, Yue and Mishra, Bhavana Dalvi and Wang, Chi and Tang, Xiangru and Zhao, Jieyu and Liu, Ge and Pfister, Tomas and Yoon, Jinsung},
  journal={arXiv preprint arXiv:2605.29054},
  year={2026}
}