Machine learning researchers often face a frustrating challenge: the code they write is tied to specific hardware and software ecosystems, limiting where and how their models can run. For example, many AI projects are built using PyTorch, a popular framework optimized for NVIDIA GPUs, while others use JAX, which works best on Google’s TPUs. Moving entire code repositories between these two environments is complex, time-consuming, and prone to errors. A new study introduces XStack, a benchmark designed to measure how well automated tools can migrate machine learning code from PyTorch on GPUs to JAX on TPUs, potentially freeing researchers to choose the best hardware for their needs.
Key Takeaways
- XStack is the first benchmark focused on migrating full machine learning repositories between different software frameworks and hardware platforms.
- It tests automated agents by giving them a PyTorch codebase and a clean TPU environment, then evaluating if the converted JAX code reproduces the original results.
- Current AI agents show limited success: the best system, Claude Opus 5, passes 75% of tests, while others like GPT-5.6-sol reach 45%, and most fall below 40%.
- This benchmark highlights the difficulty of reliable code migration and sets a measurable goal for improving cross-platform compatibility in AI research.
At its core, the research tackles a practical and technical problem: machine learning code is often “locked” into the environment where it was first developed. PyTorch and JAX are two leading frameworks that not only use different programming styles but also target distinct hardware with unique capabilities. Migrating code between them means more than just translating syntax—it requires understanding how computations map onto GPUs versus TPUs, managing dependencies, and ensuring that the final results match the original model’s performance. This process is long and complicated, especially when dealing with entire repositories rather than isolated scripts.
The team behind XStack created a set of tests that simulate this migration task. They provide automated agents with a PyTorch project and a fresh TPU setup, challenging the agents to produce a JAX version of the code that runs correctly and achieves the same published results. The benchmark includes a grading system that authenticates and calibrates outputs against the original repository’s outcomes, ensuring a fair and objective evaluation.
To put it simply, an “agent” here is an AI-powered tool designed to rewrite and adapt code across frameworks and hardware. The researchers evaluated several such agents, including Claude Opus 5 and GPT-5.6-sol, to see how well they handle the migration. While Claude Opus 5 showed promising results with a 75% pass rate, most other agents struggled to reliably convert the code, highlighting that current technology still has significant room for improvement.
By releasing XStack publicly, the researchers aim to provide the AI community with a standardized way to measure and improve the reliability of machine learning code migration. This could help developers and scientists break free from ecosystem lock-in, allowing them to leverage the best hardware available without rewriting their projects from scratch. In the future, advances spurred by benchmarks like XStack may reduce the manual effort involved in adapting ML code, accelerating research and deployment across diverse platforms.
Based on research published on arXiv by Linrong Cai, Taiming Lu, Haozhe Chen et al..