XStack: A Benchmark for ML Repo Migration

Overview§

XStack is a benchmark for evaluating AI agents on the task of migrating machine learning repositories between two dominant ecosystems: PyTorch on NVIDIA GPUs and JAX on Google TPUs. The benchmark addresses a practical need: researchers often want to run ML code on different hardware than it was originally written for. Migration is non-trivial because it requires translating code across frameworks (PyTorch → JAX) and adapting to different hardware backends (GPU → TPU), while preserving the original model's performance.

Task Definition§

Each agent is given:

• A source repository written in PyTorch, designed to run on NVIDIA GPUs.
• A clean TPU machine.

The agent must produce a JAX repository that runs on the TPU and meets a grading criterion. The criterion is based on the source repository's published result (e.g., accuracy on a benchmark), which the authors authenticate and calibrate. An agent's output is considered successful only if the migrated repository achieves a result within a tolerance of the source's published performance.

Evaluation Metrics§

Success is binary per repository: pass if the migrated code reproduces the published result within tolerance. The benchmark reports pass rates across a set of repositories (likely varying in complexity).

Results§

The paper reports pass rates for various agents:

• Claude Opus 5: 75%
• GPT-5.6-sol: 45%
• All other agents: at most 40%

These results indicate that even the best agent fails one in four migrations, and most agents perform poorly, underscoring the difficulty and unreliability of current agents for long-horizon, cross-framework, cross-hardware code migration.

Implementation Details§

The benchmark likely includes a diverse set of repositories covering different ML domains (vision, NLP, etc.) and complexities. The grading process involves running the migrated code on a TPU, measuring the target metric, and comparing to the source's published result. Authentication and calibration ensure fairness and reproducibility.

Challenges§

Migration involves:

• API mapping: translating PyTorch operations to JAX equivalents (e.g., torch.nn.Module to flax.linen.Module, torch.optim to optax).
• Device management: adapting to TPU-specific parallelism and memory constraints.
• Numerical differences: ensuring gradients and losses match despite different default precisions and random seeds.
• Dependency and build system changes.

Example: Migrating a Simple Model§

A typical migration step involves converting a PyTorch training loop to JAX. Below is a simplified illustration of the transformation.

PyTorch (source):

import torch
import torch.nn as nn
import torch.optim as optim

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(784, 10)
    def forward(self, x):
        return self.fc(x)

model = Net()
optimizer = optim.SGD(model.parameters(), lr=0.1)
criterion = nn.CrossEntropyLoss()
for x, y in dataloader:
    optimizer.zero_grad()
    loss = criterion(model(x), y)
    loss.backward()
    optimizer.step()

JAX (target):

import jax
import jax.numpy as jnp
from flax import linen as nn
import optax

class Net(nn.Module):
    @nn.compact
    def __call__(self, x):
        return nn.Dense(10)(x)

model = Net()
params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 784)))
optimizer = optax.sgd(0.1)
opt_state = optimizer.init(params)

@jax.jit
def train_step(params, opt_state, x, y):
    def loss_fn(params):
        logits = model.apply(params, x)
        return optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()
    loss, grads = jax.value_and_grad(loss_fn)(params)
    updates, opt_state = optimizer.update(grads, opt_state)
    params = optax.apply_updates(params, updates)
    return params, opt_state, loss

for x, y in dataloader:
    params, opt_state, loss = train_step(params, opt_state, x, y)

This transformation requires understanding both frameworks' paradigms: PyTorch's object-oriented, mutable state vs. JAX's functional, pure functions with explicit state.

Conclusion§

XStack provides a measurable benchmark for cross-framework, cross-hardware ML repository migration, revealing that current agents are not yet reliable for this task. It aims to spur progress toward freeing ML code from its ecosystem.