🤖 AI

A Multi-agent AI System for Deep Learning Model Migration from TensorFlow to JAX

This paper presents a multi-agent AI system that automates the migration of deep learning models from TensorFlow to JAX by combining static analysis with AI planning and coding agents, achieving a 6.4x to 8x speedup while ensuring code quality through AI-driven evaluation in large-scale commercial environments.

Original authors: Stoyan Nikolov, Bernhard Konrad, Moritz Gronbach, Niket Kumar, Ann Yan, Varun Singh, Yaning Liang, Parthasarathy Ranganathan

Published 2026-03-31

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Stoyan Nikolov, Bernhard Konrad, Moritz Gronbach, Niket Kumar, Ann Yan, Varun Singh, Yaning Liang, Parthasarathy Ranganathan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a massive, high-tech restaurant that has been cooking delicious meals for years using a very old, complex recipe book called TensorFlow. The recipes work, but the kitchen is changing. The new head chef wants everyone to switch to a modern, faster, and more efficient kitchen system called JAX.

The problem? The old recipes are thousands of pages long, written in a language the new kitchen doesn't speak, and they rely on specific, weird tools that don't exist in the new kitchen. If you tried to rewrite every single recipe by hand, it would take your team of expert chefs hundreds of years.

This paper describes how Google built a team of AI robots to do this rewriting job for them, turning the old recipes into new ones in a fraction of the time.

Here is how their "AI Kitchen Crew" works, broken down into simple parts:

1. The Problem: Why Humans Can't Do It Alone

Think of the old recipes (TensorFlow) as being written in a language where the chef remembers everything they did while cooking (stateful). The new kitchen (JAX) requires the chef to write down every single step explicitly and forget everything once the dish is served (stateless).

The Challenge: It's like asking a human to translate a 5,000-page novel from a language where grammar changes based on the mood of the sentence, into a language where grammar is strictly mathematical. If a human tries to do this alone, they will get tired, forget a page, or invent a new word that doesn't exist.

2. The Solution: The AI Multi-Agent Team

Instead of one super-smart robot trying to do everything at once (which often leads to confusion and hallucinations), the authors built a team of specialized robots, each with a specific job.

🧠 The Planner (The Architect)

Imagine a master architect who looks at the old building and draws a blueprint for the new one.

What they do: The Planner doesn't write any code. Instead, it reads the old recipe, figures out all the dependencies (which ingredients need which tools), and breaks the massive task into tiny, manageable chunks.
The Analogy: It's like saying, "First, we need to build the foundation. Then, we build the kitchen walls. Then, we install the stove." It ensures the team doesn't try to paint the roof before the walls are up.

🎻 The Orchestrator (The Conductor)

Think of a conductor in an orchestra. They don't play the instruments; they tell the musicians when to start and stop.

What they do: The Orchestrator takes the Architect's blueprint and hands out one small "chunk" of work to the coder at a time. It makes sure the coder has the right instructions and doesn't get overwhelmed by too much information. If the coder gets stuck, the Conductor decides whether to try again, skip the step, or ask for help.

📝 The Coder (The Builder)

This is the worker bee who actually writes the new code.

What they do: They take the specific instruction (e.g., "Convert this specific ingredient list") and write the new JAX code.
The Safety Net: Before they say "I'm done," they have to build the dish and taste-test it. If it fails, they fix it immediately. They also write a summary note of what they changed so the next robot knows what happened.

3. The Secret Sauce: The "Playbooks"

How do these robots know exactly how Google wants the code written? They don't just guess. They use Playbooks.

The Analogy: Imagine a set of rulebooks.
- General Rules: "Always use a fork, never a spoon."
- Style Rules: "If you write a list, use bullets, not numbers."
- Client-Specific Rules: "For the YouTube menu, always use spicy sauce."
The Magic: The team created these playbooks by looking at a few examples of recipes that humans had already successfully rewritten. They used AI to study those examples and create a "cheat sheet" for the robots. This ensures the new code looks and feels exactly like the rest of the restaurant's menu.

4. The Quality Control: The "Judge"

How do you know the new recipes are good if you don't have the original taste testers?

The Problem: You can't just run the old and new code side-by-side to see if they taste the same because the ingredients are so different.
The Solution: They created a Checklist Judge. This is another AI that doesn't look at the old code at all. It just looks at the new code and asks: "Does this dish have a sauce? Does it have a garnish? Is the temperature right?"
The Result: It gives a score based on a strict checklist of requirements, ensuring the new code is complete and follows the rules, even without a perfect "taste test."

5. The Results: Speed and Quality

The results were incredible:

Speed: The AI team finished the migration 6 to 8 times faster than a team of human experts.
Quality: For complex, huge models, the AI team was much better at remembering all the tiny details (like specific math formulas or error handling) that humans tend to miss when they are tired.
The "Virtuous Circle": The best part is that the AI is helping to build the very tools it uses. By automating the migration, they are making the future of AI development faster and easier.

In a Nutshell

This paper is about taking a massive, messy, manual job (rewriting thousands of lines of complex code) and solving it by building a team of specialized AI robots that talk to each other, follow strict rulebooks, and check their own work. It's not just about one smart robot; it's about a well-organized factory line that turns old, clunky code into modern, efficient code, saving companies years of work.

1. Problem Statement

Google faces a critical engineering challenge: migrating thousands of production deep learning models from TensorFlow (a stateful, graph-based framework) to JAX (a stateless, functional framework optimized for TPUs).

Architectural Divergence: The two frameworks differ fundamentally. TensorFlow uses implicit state management (e.g., BatchNorm updates via flags) and object-oriented layers, whereas JAX requires explicit state handling (passing batch_stats manually) and functional composition.
Complexity: Production models often rely on custom operations, complex dependency graphs, and strict internal coding standards.
Limitations of Current AI: Generic Large Language Model (LLM) coding agents struggle with long-horizon tasks, often hallucinating APIs, losing context across multiple files, and failing to produce buildable, testable code without specific domain guidance.
Cost: Manual migration by human experts is prohibitively expensive, estimated to require hundreds of engineering years.

2. Methodology: The Multi-Agent System

The authors propose a bespoke multi-agent system designed specifically for the TensorFlow-to-JAX migration workflow. The system decomposes the migration into a coordinated hierarchy of roles: Planner, Orchestrator, and Coder.

A. Knowledge Injection: The Playbook Hierarchy

To address style inconsistencies and API hallucinations, the system uses a hierarchical set of Playbooks (Markdown files provided as system prompts):

General Instruction Playbook: Covers repository structure and build tools.
Style Playbook: Enforces Google-specific coding conventions (e.g., preferring nn.compact over setup()).
Task-Specific Playbook: Generated from official docs to highlight TF vs. JAX API differences.
Client-Specific Playbook: The most critical component. Generated offline by analyzing "golden examples" (human-migrated models). It uses few-shot learning to capture team-specific patterns and nuances.

B. Agent Roles

Planner Agent:
- Function: Analyzes the source code using static analysis tools (Kythe) to discover dependencies. It generates a granular, step-by-step migration plan in JSON.
- Strategy: Breaks the migration into small, buildable steps (e.g., "migrate this specific layer," "create this directory"). It ensures the plan flows from leaf nodes (independent functions) up to the main model.
Orchestrator Agent:
- Function: Acts as the bridge between the high-level plan and the Coder. It groups plan steps into manageable "sub-steps" to optimize context windows.
- Dynamics: It performs dynamic knowledge injection (selecting relevant playbooks) and manages failure handling (retry, skip, or abort). It persists progress to a shared memory bank.
Coder Agent:
- Function: A ReAct-based agent that executes specific sub-steps. It has access to tools for reading/writing files, searching code, running builds, and executing unit tests.
- Self-Correction: The agent is instructed to build and test code after every step. If a build fails, it must fix it before proceeding. It also generates a "summary" of changes to maintain project-level context for subsequent steps.

C. Shared Memory Bank

Agents communicate via a file-based memory bank rather than direct handoffs. This stores:

Migration state and versioned snapshots.
Debugging artifacts.
Coder Summaries: A markdown summary of changes made in previous steps, ensuring the Coder has a global overview and does not undo previous work.

3. Evaluation Methodology

The system was evaluated on six real-world, high-complexity YouTube models (ranging from 0.5k to 2.9k lines of code, with up to 400 metrics).

Quality Metric (The "Judge"): Instead of relying solely on numerical equivalence (which is fragile due to framework differences), the authors developed a Checklist-based LLM Judge.
- Human experts create a "Golden Checklist" of framework-agnostic requirements (e.g., "User embeddings are L2 normalized," "Point-wise Attention layer exists").
- A dedicated LLM agent (the Judge) inspects the migrated JAX code only (blind audit) to verify if checklist items are met.
Baseline Comparisons: The system was tested against:
- Single-agent baselines (with and without playbooks).
- Multi-agent systems (with and without client-specific playbooks).

4. Key Results

The study demonstrates that the full multi-agent system with client-specific playbooks (multi_agent_yt_specific) significantly outperforms all other configurations.

Quality Improvement:
- The full system achieved checklist scores of 0.65–0.95 on complex models.
- Single-agent baselines scored significantly lower (0.09–0.66), often failing to include loss functions, metrics, or key logic.
- Statistical analysis (paired t-test, $p < 0.02$ ) confirmed the superiority of the multi-agent approach with client-specific knowledge.
Speedup:
- For complex production models, the system achieved a 6.4x to 8x speedup compared to human-only migration.
- While small models saw a modest 1.7x speedup (limited by human review time), the efficiency gains scale dramatically with model complexity.
Complexity Dependency: The performance gap between the multi-agent system and single agents widened as model complexity increased, proving the system's value for large-scale migrations.

5. Significance and Contributions

Novel Multi-Agent Architecture: The paper validates that decomposing complex code transformation into Planning, Orchestration, and Coding roles, supported by a shared memory bank, solves the context-loss and hallucination problems inherent in single-agent LLMs.
Client-Specific Knowledge Distillation: It demonstrates a scalable method for generating Client-Specific Playbooks from a small number of "golden examples," allowing the system to adapt to specific organizational styles and constraints without manual rule-writing.
Virtuous Cycle of AI Development: The system effectively uses AI to improve the development workflow of AI infrastructure itself (migrating ML frameworks), creating a self-reinforcing loop of productivity.
Practical Evaluation Framework: The introduction of Checklist-based LLM Judges offers a robust alternative to numerical equivalence testing for evaluating complex, real-world code migrations where ground-truth parity is difficult to establish.

Conclusion

This research provides a blueprint for automating high-stakes, large-scale software engineering tasks. By combining static analysis, hierarchical multi-agent collaboration, and domain-specific knowledge distillation, the system successfully bridges the gap between legacy and modern ML frameworks, reducing migration time by an order of magnitude while maintaining high architectural fidelity.