Back to News

Google Researchers Introduce T2J-Bench for PyTorch-to-JAX Migration

#coding-agents#pytorch-to-jax#benchmark#twinharness

Google researchers and collaborators introduced T2J-Bench, a benchmark of 63 repository-migration tasks for evaluating coding agents on execution-grounded PyTorch-to-JAX migration, along with the TwinHarness method. In tests across 20 agent-model configurations, the strongest baseline scored 0.676; TwinHarness improved Claude-Code with opus-5 from 0.649 to 0.723, a relative gain of 11.4%.

Coverage timeline

  1. 机器之心机器之心

    从权重转换、模型计算到训练循环和推理服务,T2J-Bench 用 63 个整仓任务检验 PyTorch 到 JAX 的迁移质量。20 组现成 agent–model 配置中,最强基线得分为 0.676 ;保持 Claude-Code 与 opus-5 不变,TwinHarness 将得分 从 0.649 提升至 0.723,相对提升 11.4% 。 图 1 | 论文首页的题名、作者与机构信息截图。 论文题目:T2J-Bench: Benchmarking Coding Agents on Execution-Grounded PyTorch-to-JAX Repository Migration 论文标题:T2J-Bench:在基于执行的 PyTorch 到 JAX 迁移库上对编码代理进行性能评估 论文:https://arxiv.org/abs/2605.29054 项目主页:https://ft2023.github.io/T2j-web/ 开源代码:https://github.com/google-research/t2j-bench/ 把一个 PyTorch 模型仓库迁到 JAX,需要协同重写模型定义、权重加载、训练循环和推理缓存。一次看似很小的修改,也可能改变梯度方向或后续训练轨迹。把这项工程交给 Coding Agent,究竟能做到什么程度? 来自 Google Cloud AI Research、伊利诺伊大学厄巴纳香槟分校、南加州大学、圣母大学和 Google DeepMind 的团队提出 T2J-Bench,并设计 TwinHarness,让 agent 利用可运行的原始仓库,逐轮定位、修复并验证行为差异。研究关注的核心是: 如何把源代码已经提供的执行证据,转化为有效的迁移反馈。 整仓迁移需要保留哪些行为? PyTorch 常以即时执行、有状态模块和面向对象的方式组织计算;JAX 则强调显式参数、函数式表达和可编译执行。迁移过程中,agent 需要重新组织程序的执行方式,让模型、训练状态与运行入口在新的框架中协同工作。 论文给出一个直观例子:先对误差平方再求平均,与先求平均再平方,只差一处计算顺序,却会得到不同的损失和梯度。代码可能正常运行,训练过程却已偏离原始实现。类似问题也会出现在优化器更新、注意力掩码以及推理阶段的 KV cache 中。 图

Related stories

  • Google open-sources Harness², a recursive self-improvement system for LLMs

    Google has open-sourced Harness², a recursive self-improvement (RSI) system that enables LLMs to modify their own code, prompts, tools, and skill libraries during runtime. The system uses a dual-layer recursive architecture—task-level and harnessing-level—to address the lack of ground truth and overfitting in real-world tasks. In tests, equipping GLM-5 with Harness² improved its score on the WorkBuddy-Bench office benchmark by 11.5 points, surpassing native Gemini 3.5 Flash. The project is fully open-sourced and supports OpenCode and Codex.

    1 reports · 1 sources
  • DeepMind Researchers Find Numerical Precision Bugs in Softmax and Cross-Entropy Across Major Frameworks

    Google DeepMind researchers identified two numerical precision bugs in the standard Softmax and cross-entropy implementations used by PyTorch, JAX, and TensorFlow. The bugs can subtly alter gradients during training, potentially causing loss spikes or capability regression. The team proposed a fix that adjusts the computation without increasing floating-point precision or redesigning the loss function; the paper was submitted to arXiv on October 22.

    1 reports · 1 sources
  • Fudan, Tencent Hunyuan, Tsinghua Release ExplorationBench to Test AI Rule Discovery

    Fudan University, Tencent Hunyuan, and Tsinghua University released ExplorationBench, a benchmark measuring how AI systems discover rules in unfamiliar environments. It includes two altered 'alien worlds'—the programming language AlienCode and the logic proof system AlienLogic—where models must learn rules through trial and error. Results showed that without environment feedback, scores remained low (max 11.0%), while autonomous exploration led to high scores (up to 89.0%).

    2 reports · 1 sources
  • Google launches TorchTPU, vLLM adds TPU backend

    Google announced TorchTPU, an open-source version of PyTorch that lets AI applications run on TPUs and Google Cloud via simple configuration changes without rewriting code. Concurrently, vLLM introduced vllm-torchtpu, a new backend for Google Cloud TPUs built on TorchTPU, featuring Pallas kernels, multi-host parallelism, KV cache offloading, and disaggregated serving. These moves, following Amazon's similar tool, challenge Nvidia's dominance in AI development.

    3 reports · 3 sources

News updates every hour. Subscribe in your RSS reader to get new stories automatically.

Subscribe