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

summary

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

In short

The paper introduces T2J-Bench, a benchmark for migrating PyTorch code to JAX by framing conversion as transfer under a fixed equivalence contract. It tests coding agents on whole-repository migrations across 63 tasks. Findings show that existing benchmarks fail because they don't require actual behavioral equivalence, and current agents overestimate success significantly when validated against the fixed evaluator.

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 used across episodes

Reinforcement Learning (RL) — explained in 25 episodes
Sample Complexity — explained in 18 episodes
Federated Learning — explained in 17 episodes
Retrieval-Augmented Generation (RAG) — explained in 16 episodes
Interpretability — explained in 15 episodes
Quantum Fisher Information (QFI) — explained in 15 episodes

This episode discusses

The paper

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

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

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.

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.

More episodes

← Home