XStack Benchmark: A Touchstone for Cross-Framework ML Code Migration

Title

XStack Benchmark: A Touchstone for Cross-Framework ML Code Migration

Meta Description

XStack is a benchmark for evaluating agents' ability to migrate ML codebases between the PyTorch and JAX ecosystems. Claude Opus 5 achieves a 75% pass rate, GPT-5.6-sol 45%, and all other agents at most 40%. Learn the details and significance of this cross-framework migration challenge.

Article Content

Machine learning code should run where the researcher chooses, not where it was originally written. Today, ML code primarily inhabits two ecosystems: PyTorch with NVIDIA GPUs, and JAX with Google TPUs. Migrating a codebase between these two ecosystems means building and testing across frameworks and hardware. This migration is both a practical need—it unlocks different hardware for researchers—and a natural benchmark: it involves different frameworks and hardware, covers the entire codebase, and is long-horizon and challenging. To this end, we introduce XStack, a benchmark for evaluating agents' ability to migrate ML codebases between the two major ecosystems. Each agent receives a PyTorch codebase and a clean TPU machine, and must deliver a JAX codebase. Its output must meet our scoring criteria based on the source repository's published results (certified and calibrated by us) to pass. Currently, agent migration remains unreliable: Claude Opus 5 achieves a 75% pass rate, GPT-5.6-sol 45%, and all other agents at most 40%. XStack is a step toward liberating ML code from its ecosystem constraints. We release XStack to make reliable migration across frameworks and hardware measurable and to drive progress.

Why Do We Need the XStack Benchmark?

Ecosystem lock-in of ML code is a real problem. Researchers choose PyTorch or JAX often constrained by hardware availability rather than personal preference. Migrating a codebase requires not only understanding the API differences between the two frameworks, but also handling hardware characteristics, performance tuning, and the dependencies of the entire repository. This complexity makes migration an ideal testbed for evaluating agent capabilities: it demands long-horizon planning, cross-domain knowledge integration, and a global grasp of the entire codebase. XStack formalizes this challenge, providing agents with a standardized arena.

Core differences between JAX and PyTorch: PyTorch uses eager execution mode, where operations are computed immediately upon invocation, making debugging intuitive; JAX is based on the functional programming paradigm, using just-in-time (JIT) compilation by default to compile pure functions into XLA computation graphs, requiring side-effect-free code and static shapes. Additionally, JAX's automatic differentiation is implemented through the grad transform rather than PyTorch's .backward() method; parameter updates must be done explicitly with libraries like optax rather than an optimizer's .step(). These paradigm differences mean migration is not simply API replacement, but requires restructuring the code's execution model and state management.

TPU and the XLA compiler: TPU is Google's application-specific integrated circuit (ASIC) designed specifically for tensor computation, and its execution relies on the XLA (Accelerated Linear Algebra) compiler to compile computation graphs into TPU instructions. Unlike the CUDA kernels commonly used on GPUs, XLA performs global optimizations such as operator fusion and memory planning, but requires static computation graphs with known shapes. This means that when running on TPU, dynamic control flow and variable-shape operations require special handling (e.g., jax.lax.cond, jax.lax.scan), which is also one of the main performance and compatibility challenges when migrating PyTorch code to JAX+TPU.

How Does XStack Work?

The core design of XStack is simple yet strict. Each agent receives a PyTorch repository and a clean TPU machine, with the task of producing a JAX repository. Scoring does not rely on subjective judgment but is based on the source repository's published results—certified and calibrated to ensure the migrated code is functionally equivalent and meets performance targets. This real-results-based scoring mechanism makes the benchmark both close to practice and reproducible. Agents must independently complete the entire migration process: from code parsing and framework conversion to hardware adaptation and test validation.

Agent Performance: A Clear Gap

Preliminary evaluation reveals uneven agent performance on the migration task. Claude Opus 5 leads with a 75% pass rate, GPT-5.6-sol reaches 45%, and all other agents do not exceed 40%. This distribution shows that although top models demonstrate some potential, migration reliability is far from production-grade. The reasons behind the high pass rates may include stronger code comprehension and long-context handling capabilities, but even the leader has a quarter of failure cases, indicating the task still presents significant challenges.

Agents and pass rates in the benchmark: Claude Opus 5, GPT-5.6-sol, and others mentioned in the article are AI agents (LLM-based code agents) participating in the evaluation. Pass rate refers to the proportion of XStack tasks in which an agent successfully delivers a JAX codebase meeting the scoring criteria. A 75% pass rate means functionality and performance meet targets on most tasks, but a quarter of failure cases remain; the at-most-40% pass rates of other agents reflect the general limitations of current models on long-horizon, cross-framework migration tasks.

Migration Challenges and Future Directions

The challenges of migration are multidimensional. Semantic differences between frameworks (such as automatic differentiation mechanisms and tensor operation conventions), hardware-specific optimizations (such as TPU's XLA compilation and GPU's CUDA kernels), and repository-level dependency management all constitute obstacles. The release of XStack provides the community with a measurable goal: by quantifying agents' migration capabilities, it drives progress in cross-framework tools and methods. In the future, with iterations of the benchmark and improvements in agent capabilities, we expect to see more reliable migration solutions, ultimately allowing ML code to flow freely between different ecosystems.

One-line Summary

XStack is a benchmark evaluating agents' ability to migrate ML codebases between PyTorch and JAX/TPU.

Paragraph Summary

XStack is a benchmark for evaluating AI agents' ability to migrate machine learning codebases across frameworks and hardware. Each agent receives a PyTorch repository and a clean TPU machine and must deliver a functionally equivalent JAX codebase that meets performance targets, with scoring based on the source repository's certified published results. Preliminary evaluation shows agent migration remains unreliable: Claude Opus 5 achieves a 75% pass rate, GPT-5.6-sol 45%, and other agents at most 40%. Migration challenges involve framework semantic differences, hardware-specific optimizations, and repository-level dependency management. XStack aims to make cross-ecosystem migration measurable and to drive ML code to break free from ecosystem constraints.