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

Stoyan Nikolov; Bernhard Konrad; Moritz Gronbach; Niket Kumar; Ann Yan; Varun Singh; Yaning Liang; Parthasarathy Ranganathan

cs.SEMar 28, 2026

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

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

Organizations: Google Munich, Germany · Google Mountain View, California, USA · Google Sunnyvale, California, USA

Abstract

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.

Model	Line Count (Main/Deps.)	Layers	Metrics
small_1	0.5k / $\sim$ 500	$\sim$ 20	$\sim$ 20
small_2	0.4k / $\sim$ 500	$\sim$ 17	$\sim$ 25
medium_1	2.2k / $\sim$ 1k	$\sim$ 40	$\sim$ 30
medium_2	2.3k / $\sim$ 3k	$\sim$ 40	$\sim$ 60
large_1	2.7k / $\sim$ 13k	$\sim$ 250	$\sim$ 400
large_2	2.9k / $\sim$ 9k	$\sim$ 130	$\sim$ 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	$\sim$ 16 hrs	$\sim$ 2 hrs	8x
model_2	4k / 20k	250	300	$\sim$ 32 hrs	$\sim$ 5 hrs	6.4x

Table 5. Migration Efficiency and Human-in-the-Loop Metrics

AI Coding Agents

AI coding agents use language models and software tools to plan, edit, test, and maintain code with varying levels of autonomy.

See all 326 papers

Agentic Framework for Deep Learning workload migration via In-Context Learning
Jun 14, 2026Qiyue Liang, Steven Ingram, George Vanica +4Code Translation Differential Testing
Google, Sunnyvale, CA · 2Google, Kirkland, WA · 1Google, Sunnyvale, CA +2
Learning Bug Context for PyTorch-to-JAX Translation with LLMs
Oct 10, 2025Hung Phan, Son Vu, Tuan Dinh +3Code Translation LLM Code Generation
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
Design Docs Are All You Need: An AI-native Machine-Learning Performance Tool
Sep 4, 2026Samuel Kushnir, Kimia Noorbakhsh, Kavya Sreedhar +6Code Generation
Google DeepMind · MIT · Google +1

Jun 14, 2026J/K move · Enter open · S save

Agentic Framework for Deep Learning workload migration via In-Context Learning

Qiyue Liang, Steven Ingram, George Vanica, Andi Gavrilescu+3

Google, Sunnyvale, CA · 2Google, Kirkland, WA · 1Google, Sunnyvale, CA+2

Code Translation Differential Testing

Open paper AskPDF arXiv

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

Abstract

AI Coding Agents

Explore similar work

Agentic Framework for Deep Learning workload migration via In-Context Learning

Learning Bug Context for PyTorch-to-JAX Translation with LLMs

Design Docs Are All You Need: An AI-native Machine-Learning Performance Tool