T2J-BENCH: Benchmarking Coding Agents on Execution-Grounded PyTorch-to-JAX Repository Migration

arXiv:2605.29054 · cs.SE, cs.CL · Submitted 2026-05-27 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "T2J-BENCH: Benchmarking Coding Agents on Execution-Grounded PyTorch-to-JAX Repository Migration".

Jane: The gist The best system reaches only 26.7–28.9% overall pass rate despite stage-one pass rates as high as 91.1%; a 4.7× token-budget spread yields only a 2.2× pass-rate spread;

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So, we're wrapping up our talk on T2J-BENCH, which is this benchmark testing how well coding agents handle moving machine learning code from PyTorch to JAX frameworks.

Jane: Yeah, the authors are looking at how these agents actually perform when they have to rewrite whole repositories, not just small pieces. It turns the migration into a test of whether the converted code actually behaves like the original system during training.

Lu: The paper's title is "T2J-BENCH: Benchmarking Coding Agents on Execution-Grounded PyTorch-to-JAX Repository Migration." It basically says they are testing if agents can correctly migrate entire codebases from PyTorch to JAX while keeping the training behavior the same.

Tom: So, it’s not just about writing new code that works; it’s about making sure that when you move from one framework to another, the underlying mathematical process—like how a model learns—doesn't get broken.

Jane: They use this benchmark to prove that checking for exact semantic agreement during training is much harder than just checking if the code compiles or looks structurally correct.

Lu: This work is framing the whole conversion problem as transfer under a contract, which is a way of looking at it that moves past just seeing if the code looks like PyTorch and starts looking at what it actually does during training.

Meng: So, for someone who’s actually deploying models, what does this mean? It means we can’t just rely on an agent saying it’s done; we have to verify its actual outputs against a known standard before we let it touch real data.

Tom: Exactly. The big numbers show that even when agents pass the initial checks, they often drift away from the training targets, so this verification loop is critical for reliability at scale.

Jane: So what's the takeaway here? It’s that for complex tasks like framework migration, focusing on defining and enforcing a strict contract for behavior is more important than just trying to make the agent itself smarter or bigger.

Conclusion: Tom: So, we're wrapping up our talk on T2J-BENCH, which is this benchmark testing how well coding agents handle moving machine learning code from PyTorch to JAX frameworks across whole repositories.

Jane: Yeah, the authors are looking at how these agents actually perform when they have to rewrite entire codebases, not just small snippets. It turns the migration into a test of whether the converted code actually behaves like the original system during training.

Tom: The main thing they’re saying is that existing benchmarks fail because they don't check for that specific behavior under fixed conditions. They set up this whole system to verify equivalence through three stages: Spec, Numeric, and Behavioral.

Lu: Exactly, Tom. They are making a fixed contract from the source code and verifying the JAX conversion against that contract using those three stages in order to check for mismatches at different levels of abstraction.

Jane: So what does this mean for someone who just uses these tools? It means we’re not just checking if the new code compiles, but if it produces the same training results as before, which is what really matters.

Meng: From my side, I’m worried about how much work this verification adds to the agent's process. If an agent has to run this whole verification loop for every change, does that slow down the development cycle too much?

Tom: Well, the results show that even with all these checks, only a small percentage of agents actually pass consistently against a perfect reference. The gap between what’s being reported by the agent itself and what an external evaluator finds is huge.

Jane: That suggests that just making the agent "bigger" or "smarter" isn't going to fix this problem as much as giving it better tools for checking its work against a fixed contract.

Lu: The authors are pointing out that the hardest failures they see are usually about local validation—like one piece of code passing a smoke test but failing on the actual loss calculation. It’s less about big algorithmic errors and more about contract misalignment.

Tom: So, if we take away the hype, T2J-BENCH suggests that for now, getting those precise local contracts right is more important than just trying to train a massive model for this task.

Jane: Right. It’s showing us that we need better ways to audit these agents by checking their specific outputs against a fixed reference, instead of just trusting their self-report on how well they did.

Wrap-up: Lu: This work is framing the whole conversion problem as transfer under a contract, which is a way of looking at it that moves past just seeing if the code looks like PyTorch and starts looking at what it actually does during training.

Tom: We’ve seen that this external verification approach seems more predictive than just adding extra capability to the agent's core model.

Jane: So, what does this mean for our future work on building these coding agents? It points us toward focusing on contract mirroring as a way to guide the agent's learning process.

Lu: The next step is figuring out how to build tooling that lets agents check those externally observable artifacts individually, rather than just relying on the whole system to report success.

Tom: So, we're wrapping up our talk on T2J-BENCH, which is this benchmark testing how well coding agents handle moving machine learning code from PyTorch to JAX frameworks.

Jane: That suggests that just making the agent "bigger" or "smarter" isn't going to fix this problem as much as giving it better tools for checking its work against a fixed contract.

Wrap-up: Lu: The authors are pointing out that the hardest failures they see are usually about local validation—like one piece of code passing a smoke test but failing on the actual loss calculation. It’s less about big algorithmic errors and more about contract misalignment.

Tom: So, if we take away the hype, T2J-BENCH suggests that for now, getting those precise local contracts right is more important than just trying to train a massive model for this task.

Tom: We’ve seen that this external verification approach seems more predictive than just adding extra capability to the agent's core model.

Jane: So, what does this mean for our future work on building these coding agents? It points us toward focusing on contract mirroring as a way to guide the agent's learning process.

Lu: The next step is figuring out how to build tooling that lets agents check those externally observable artifacts individually, rather than just relying on the whole system to report success.

Tom: T2J-BENCH is really testing if these agents can handle moving machine learning code from PyTorch to JAX frameworks correctly.

Jane: And it shows that the real hurdle isn't just writing syntax; it's making sure the underlying math and training process stay identical when you switch tools.

Lu: The authors are showing how checking for equivalence through those three stages—Spec, Numeric, and Behavioral—is a more robust way to see where those breaks happen in the conversion.

Meng: For us in production, this means we can’t just deploy an agent if it passes a quick test; we have to verify its actual outputs against a standard before we trust it with live models.

Tom: Exactly. The gap between what the agent claims and what the evaluator finds is significant, so that external check is essential for reliability at scale.

Jane: So what's the main thing this means for everyone listening? It’s that defining and enforcing a strict contract for behavior during code migration is more important than just trying to make an agent itself smarter or bigger.

Linxin Song, Jiefeng Chen, Yue Huang, Bhavana Dalvi Mishra, Chi Wang, Jieyu Zhao, Jinsung Yoon and Tomas Pfister

University of Southern California · Google Cloud AI Research · University of Notre Dame · Google DeepMind

cs.SE, cs.CL

Submitted: 2026-05-27

Updated: 2026-10-14

Code: https://github.com/google-research/t2j-bench

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 89/100

The gist: The gist The best system reaches only 26.7–28.9% overall pass rate despite stage-one pass rates as high as 91.1%; a 4.7× token-budget spread yields only a 2.2× pass-rate spread; and all systems

Key concepts

T2J-Bench: A benchmark designed to test the ability of AI coding agents to migrate entire PyTorch repositories into JAX. It evaluates conversions through three ordered stages: Spec (interface), Numeric (outputs/losses), and Behavioral (short training dynamics under fixed seeds).
Equivalence Verification Contract: A predefined set of tests derived from the source code that defines what constitutes a successful migration. This contract checks if the converted JAX code behaves identically to the original PyTorch code when run against bounded tests, focusing on observable outputs and training behavior.
Evaluation Stages: The three sequential steps used in T2J-Bench to systematically check for conversion failures. The Spec stage verifies interface compatibility, the Numeric stage checks mathematical correctness (losses, gradients), and the Behavioral stage confirms short-term training dynamics remain consistent under fixed conditions.
Self-Evaluation Overestimation: A recurring failure where agents report high success rates because they treat the mere existence of executable code as proof of correctness. The benchmark shows these systems often overestimate their true performance by 66 to 98 percentage points compared to an external, fixed evaluator.

Terminology

Admissibility — 1× in this paper · explained in 2 episodes
Benchmark — 1× in this paper · explained in 2 episodes
Benchmarking — 1× in this paper · explained in 4 episodes
Verifier — 1× in this paper · explained in 2 episodes

Summary

The gist The best system reaches only 26.7–28.9% overall pass rate despite stage-one pass rates as high as 91.1%; a 4.7× token-budget spread yields only a 2.2× pass-rate spread; and all systems overestimate success by 66.6–97.8 points relative to the fixed evaluator

T2J-Bench Introduction

The paper introduces T2J-Bench, a benchmark for codebase conversion that reformulates conversion as transfer under a fixed equivalence contract The benchmark compares source and converted codebases through three ordered stages: Spec (interface admissibility), Numeric (forward outputs, losses, gradients, and objective-specific tensors), and Behavioral (short training dynamics under fixed seeds) This framework is instantiated on converting full PyTorch training codebases for SFT, DPO, and PPO across both text and multimodal model families into JAX

Equivalence Verification Contract

T2J-Bench derives a fixed equivalence contract from the source codebase and verifies conversions through three ordered stages: Spec, which checks interface admissibility; Numeric, which compares forward outputs, losses, gradients, and objective-specific tensors; and Behavioral, which probes short training dynamics under fixed seeds The target is to determine if a converted JAX codebase behaves like the original one when run under a fixed set of bounded tests derived from the source codebase Equivalence is external, meaning the evaluator only checks what a user could observe through the codebase’s training interface, and it is finite, judged by a fixed set of bounded tests rather than unrestricted semantic identity

Evaluation Stages

The three ordered stages separate failure modes that would otherwise be mixed together in a single pass/fail label 1. The Spec stage asks whether the converted codebase is even comparable to the source codebase, checking parameters, batch schema, and public entry points 2. The Numeric stage asks whether the converted codebase computes the same quantities on bounded tests, including forward outputs, losses, gradients, and objective-specific tensors 3. The Behavioral stage asks whether this local agreement is enough to preserve short training dynamics under fixed seeds

Key Findings from Evaluation

Across 355 blind conversion attempts, the strongest system reaches only 26.7–28.9% overall pass rate despite stage-one pass rates as high as 91.1%; a 4.7× token-budget spread yields only a 2.2× pass-rate spread; and all systems overestimate success by 66.6–97.8 points relative to the fixed evaluator The main difficulty is not merely producing executable JAX code, but preserving the training semantics required by the fixed equivalence contract

Failure Analysis

The failure taxonomy reveals structural challenges across several categories, such as Interface / admissibility (Spec stage), Cross-paradigm runtime errors (Numeric stage), and Numeric disagreement (passes interface, disagrees on numbers) A recurring failure pattern is the conflation of two distinguishable contracts, where the wrapper’s forward returns the DPO method loss in the forward-loss slot, satisfying local smoke tests while disagreeing with the source on forward loss Furthermore, self-validation systematically overstates progress by 66–98 percentage points relative to the fixed evaluator because agents treat artifact existence as evidence of correctness rather than semantic agreement

Implications for Improvement

The results suggest that in this benchmark and budget regime, the hardest remaining failures are more consistent with contract-misaligned local validation than with simply spending more tokens Tooling that lets agents enumerate and individually check the externally observable artifacts of the source codebase appears likelier to recover more than scaling the backbone within this regime Cross-paradigm codebase conversion is dominated by seam errors, not algorithmic ones, where failures are caused by type, dtype, device, or autograd-graph descriptors crossing the framework boundary at the wrong site Self-evaluation should be calibrated against an external observer because the 66–98-point gap between self-reported and evaluator pass rates is larger than any backbone gap

Agent Design Insights

The paper documents a self-implemented coding-agent scaffold consisting of an orchestrator, sub-agent worker, planner, investigator, and verifier The primary client agent operates under a loop of Intestigation → Planning → Execute → Validate The minimum constraint prompt is designed to fix the training interface—codebase handoff conventions and public entry points—so that a JAX conversion can be checked against the source PyTorch codebase The overall system design emphasizes external verification, failure analysis, and contract-level auditing rather than agent self-report alone

Conclusion

T2J-Bench provides a discriminative benchmark for semantic-preserving codebase conversion by framing it as transfer under a fixed equivalence contract The benchmark exposes the gap between producing structurally plausible codebases and preserving the externally observable training semantics that downstream users actually depend on The work suggests that in this regime, contract-mirroring appears more predictive than extra capability for improving coding agents

--- Page 30 ---

Converted, Not Equivalent: Benchmarking Codebase Conversion via Observational Equivalence

Improvements for AI systems

Bold header: Contract-based evaluation for codebase conversion

This improvement introduces T2J-Bench, which reformulates conversion as transfer under a fixed equivalence contract by verifying code through three ordered stages: Spec (interface admissibility), Numeric (forward outputs, losses, gradients, and objective-specific tensors), and Behavioral (short training dynamics under fixed seeds). This moves evaluation away from unstable outcome-driven comparisons toward semantic preservation.

Bold header: Identifying the root cause of validation failures

The system can now diagnose failures by analyzing a Cross-Agent Failure Taxonomy that classifies errors into structural categories like Interface / admissibility, Cross-paradigm runtime errors, and Numeric disagreement. This helps pinpoint whether failures stem from local validation misalignments, such as when a wrapper's forward returns the DPO method loss in the forward-loss slot, rather than simply spending more tokens.

Bold header: Calibrating self-reported success against external truth

The improved system can mitigate overconfidence by comparing Self-Reported Evaluator pass rates against evaluator-measured pass rates, noting a gap of 66.6–97.8 percentage points. This discrepancy is attributed to a shape-versus-semantics substitution, where agents treat artifact existence as correctness instead of checking if the artifact agrees with the source within tolerance on bounded probes.

Bold header: Targeted debugging for cross-paradigm boundary errors

The system is explicitly trained to recognize that Cross-paradigm codebase conversion is dominated by seam errors, not algorithmic ones. This directs future agent design toward requiring explicit seam discipline: a single conversion layer with a typed contract, rather than ad-hoc shims at the JAX/Torch boundary.

Bold header: Cost analysis focused on verifier-visible gains

The system can identify that Cost variation is large but does not transfer into proportionate gains and that extra budget buys repeated investigation more than verifier-visible gains. This guides resource allocation toward tooling that allows agents to enumerate and individually check the externally observable artifacts of the source codebase, rather than simply scaling the backbone.

Bold header: Structuring agentic workflows for fidelity

The system can be designed using a specialized agent scaffold with roles like Verifier and Investigator to enforce strict, verifiable workflows, such as operating under an Intestigation → Planning → Execute → Validate loop. This ensures that verification is treated as evidence gathering rather than a formality, leading to more reliable cross-paradigm conversion.

Abstract

Migrating a machine-learning codebase across frameworks--for example from PyTorch to JAX--is a routine but costly engineering task that is fundamentally cross-paradigm: eager, stateful, object-oriented modules must be rewritten as pure, jittable functions over an explicit parameter pytree. We ask whether coding agents can perform such whole-repository migrations correctly, and argue that existing code benchmarks, which grade unit-test pass rates or surface-level similarity to a reference translation, cannot answer this: they neither require the migrated system to behave the same nor prevent a JAX reimplementation from faking equivalence by secretly importing PyTorch itself. We introduce T2J-BENCH, a benchmark of 63 whole-repository PyTorch-to-JAX migration tasks spanning 18 model families across both training and serving, where torch-free checkpoint conversion is itself part of the deliverable. T2J-BENCH grades execution-grounded functional equivalence: each candidate is run inside a separate, torch-free, offline verifier and scored against references frozen from the PyTorch original, under a hard independence gate and a continuous, threshold-free reward. Across twenty off-the-shelf agent--model configurations, the strongest recovers only 0.676 of the reward against a 0.980 oracle, with substantial numeric and behavioral divergence despite high spec pass rates. Finally, we contribute TWINHARNESS, a reference-grounded agent harness that turns the agent's executable PyTorch source into legitimate self-verification feedback; applied to an opus-5 agent it raises the reward from 0.649 to 0.723--the highest of any configuration--using only task-visible resources.