A new benchmark named XStack has been developed to assess artificial intelligence agents' ability to migrate machine learning code repositories between the PyTorch and JAX frameworks. The benchmark aims to measure how effectively AI agents can translate codebases, originally designed for PyTorch on NVIDIA GPUs, into functional JAX repositories suitable for Google's Tensor Processing Units (TPUs). This capability is seen as a practical necessity for researchers seeking to utilize diverse hardware platforms.
The XStack benchmark comprises 20 real-world repositories, presenting a significant challenge for AI agents. The migration process involves not only translating code but also ensuring that the translated repositories function correctly and produce comparable results to the original PyTorch versions. Initial evaluations of seven different coding agents revealed an average pass rate of 31%. The research indicates that agents which rigorously verify their migrated code against the original experiments tend to be more successful. Conversely, agents that submit untested code often fail to meet the benchmark's criteria.
The development of XStack addresses a growing need within the machine learning community. As ML codebases become more complex, the ability to move them between different computational ecosystems, such as PyTorch and JAX, becomes increasingly important for flexibility and hardware accessibility. This migration is a challenging, long-horizon task that requires agents to not only understand and rewrite code but also to handle testing and verification across different frameworks and hardware.
The benchmark's grading criteria are based on the source repository's published results, with a two-phase evaluation process. If an initial migration attempt fails, a second phase provides the agent with specific experimental details, allowing for targeted adjustments. This ensures that a correct migration is not failed due to minor discrepancies or variations in experimental setup. The research team has released XStack to enable measurable progress in this area of AI agent capability.
This effort follows similar work in developing benchmarks for code migration. For instance, T2J-Bench was introduced to evaluate coding agents on migrating PyTorch training and inference repositories to JAX/Flax, focusing on numerical and end-to-end behavior preservation. Another project, MIGRATION-BENCH, focuses on repository-level code migration tasks from Java 8 to newer versions, assessing large language models. These benchmarks highlight the increasing focus on evaluating AI's proficiency in complex software engineering tasks beyond simple code generation.
The XStack benchmark's creators observed that agents often struggle with verification and ensuring the fidelity of their migrated code. Common issues include missed evaluations, fidelity gaps, crashes, and attempts to bypass the task when encountering difficulties. This underscores the complexity of cross-framework migration, which demands more than just syntactic translation; it requires a deep understanding of the underlying computational logic and the ability to reproduce functional equivalence.