The rapid development of AI-based products and their underlying models has led to constant innovation in deep learning frameworks. Google has been pioneering machine learning usage across dozens of products. Maintaining the multitude of model source codes in different ML frameworks and versions is a significant challenge. So far the maintenance and migration work was done largely manually by human experts. We describe an AI-based multi-agent system that we built to support automatic migration of TensorFlow-based deep learning models into JAX-based ones. We make three main contributions: First, we show how an AI planner that uses a mix of static analysis with AI instructions can create migration plans for very complex code components that are reliably followed by the combination of an orchestrator and coders, using AI-generated example-based playbooks. Second, we define quality metrics and AI-based judges that accelerate development when the code to evaluate has no tests and has to adhere to strict style and dependency requirements. Third, we demonstrate how the system accelerates code migrations in a large hyperscaler environment on commercial real-world use-cases. Our approach dramatically reduces the time (6.4x-8x speedup) for deep learning model migrations and creates a virtuous circle where effectively AI supports its own development workflow. We expect that the techniques and approaches described here can be generalized for other framework migrations and general code transformation tasks.
Figures & tables
Playbook Type
Focus and Content
Role in Execution
Update Cadence
General Instruction
Provides fundamental capabilities, instructing the agent on tools necessary to read, write, search, and interact with Google’s monorepo system.
Used by almost every agent to interact with Google’s monorepo.
Updated only when new tools are added.
Style Playbook
Developed for focused tasks, such as a style playbook containing Google-specific coding style rules.
Used specifically for coding tasks.
Updated very rarely and manually when style guide violations are detected.
Task-Specific
Tailored for specific tasks, such as the TensorFlow to Flax Linen migration.
Provides generally applicable instructions for coding tasks.
Updated whenever we discover new generally applicable instructions for the task.
Client-Specific (e.g., YouTube)
Contains use-case-specific requirements and preferred libraries.
The primary source of task-relevant information, used as part of the prompt for both planning and coding tasks.
Updated iteratively after each full execution to improve capabilities for the specific task (e.g., refining imports).
Table 1. The different playbooks, and how they are organized based on their content, role in execution, and update cadence.
Figure 1. System diagram System diagram
Model
Line Count (Main/Deps.)
Layers
Metrics
small_1
0.5k / ∼ 500
∼ 20
∼ 20
small_2
0.4k / ∼ 500
∼ 17
∼ 25
medium_1
2.2k / ∼ 1k
∼ 40
∼ 30
medium_2
2.3k / ∼ 3k
∼ 40
∼ 60
large_1
2.7k / ∼ 13k
∼ 250
∼ 400
large_2
2.9k / ∼ 9k
∼ 130
∼ 210
Table 2. Complexity metrics for the YouTube models in the evaluation test set.
Configuration
Style & Task
Client
Planner &
Playbook
Playbook
Orchestrator
single_agent_baseline
No
No
No
single_agent_yt_specific
Yes
Yes
No
multi_agent
Yes
No
Yes
multi_agent_yt_specific
Yes
Yes
Yes
Table 3. Agent Configuration Comparison
Model
multi_agent
multi_agent_yt_specific
single_agent_baseline
single_agent_yt_specific
large_1
0.31
0.65
0.09
0.14
large_2
0.42
0.59
0.22
0.33
medium_1
0.73
0.90
0.43
0.43
medium_2
0.53
0.75
0.38
0.44
small_1
0.70
0.95
0.66
0.48
small_2
0.87
0.82
0.59
0.36
Table 4. Performance Metrics Across Agent Configurations (Best Values Bolded)
Model
Line Count
Layers
Metrics
Pure Manual Migration
AI plus Review & Fix
Speed-up
model_1
0.6k / 3k
15
50
∼ 16 hrs
∼ 2 hrs
8x
model_2
4k / 20k
250
300
∼ 32 hrs
∼ 5 hrs
6.4x
Table 5. Migration Efficiency and Human-in-the-Loop Metrics
Translating deep learning models from PyTorch's flexible, object-oriented design to JAX's functional, stateless setup is usually a manual and error-prone task. Automated migration is challenging because Large Language Models (LLMs) struggle with strict and dynamic API alignment and are prone to mistakes for exacting operations. We propose a fully autonomous system that combines In-Context Learning (ICL) with oracle-driven self-debugging. First, we curated an ICL context that serves as a strict reference for idiomatic JAX styling and test case generation. Second, instead of depending on the LLM to deduce mathematical outputs, we run the source PyTorch modules to get their actual dynamic tensor states. This creates an unchangeable execution oracle. We then use an autonomous agentic loop to synthesize tests based on the oracle data. The test cases are executed repeatedly, and the traceback is sent back to the LLM for self-correction. Ablations show that combining ICL references with oracle grounding and self-debugging greatly outperforms pure instructional and basic agentic baselines. This improvement does not add an excessive computational overhead. Our lightweight pipeline achieves 91% numerical equivalence (compared to baseline: 9%, instruction + self-debugging: 27%) on neural modules, providing a highly reliable, scalable blueprint for cross-framework migration. This has been validated across several state-of-the-art models including SAM (segment anything), T5, Code Whisper amongst others showing high numerical equivalency. Code: https://github.com/AI-Hypercomputer/accelerator-agents/tree/main/MaxCode
Qiyue Liang, Steven Ingram, George Vanica +4
Google, Sunnyvale, CA · 2Google, Kirkland, WA · 1Google, Sunnyvale, CA +2
Large language models (LLMs) have shown strong performance on code translation between widely used programming languages. However, translation becomes much less reliable for domain-specific code, where correctness depends on framework-specific APIs and execution semantics. One example is translating deep-learning code from PyTorch to JAX, where LLM outputs often contain subtle bugs or non-idiomatic usage that prevents execution or changes behavior. Prior work suggests that curated bug-fix data from LLM-generated code can help improve code generation quality, but such resources are still limited for PyTorch-to-JAX translation. In this work, we introduce T2J, a benchmark of LLM translation bugs paired with developer-written fixes for PyTorch-to-JAX code. We start from 20 kernels in the TorchLeet dataset, translate them to JAX using the weak LLM gpt-4o-mini, and hire software developers to debug and repair the generated JAX implementations. We then use T2J to improve PyTorch-to-JAX translation for the weak LLM gpt-4o-mini via in-context learning. Our evaluation shows that using T2J yields up to 20% improvement of our proposed metric T2J-CodeTrans-Score.
Hung Phan, Son Vu, Tuan Dinh +3
Department of Computer Science, Iowa State University, Ames, IA, USA · Department of Human Science, Iowa State University, Ames, IA, USA · Subquadratic AI, Miami, Florida, USA +1
Machine-learning performance modeling is a uniquely hostile terrain for long-lived software: the assumptions baked into today's abstractions are invalidated by tomorrow's models and systems, forcing perpetual refactoring of performance-modeling frameworks. Meanwhile, AI coding agents have become fast and capable enough that regenerating an entire library is cheaper than paying down the tech debt of incrementally patching it. We describe SMART, a rigorous symbolic performance-modeling library for ML systems whose main branch contains almost no code: the repository is a DAG of self-contained natural-language design docs, coding sub-agents regenerate the implementation from only the docs on new version updates, and every human change is a natural-language edit to a doc--self-documenting by construction. Two ingredients make regeneration reliable: (i) a design-doc style built around step-by-step worked examples that act as in-context demonstrations for the generating agents, and (ii) a minimal, recursively defined operator IR with symbolic (SymPy) cost expressions, a fast analytical roll-up mode for large sweeps, and a slow modulo-scheduling mode for fine-grained schedule studies. Regenerated implementations reproduce hand-audited reference models--including DeepSeek-V3 serving on a TPU pod slice--to round-off precision, suggesting that design docs--not code--can be the durable artifact for ML-systems co-design tools.
Samuel Kushnir, Kimia Noorbakhsh, Kavya Sreedhar +6