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

2026-03-28

AI Summary

This work addresses the inefficiency and high cost of migrating deep learning models across frameworks—such as from TensorFlow to JAX—in large-scale AI systems. To tackle this challenge, the authors propose an automated, multi-agent collaborative migration approach that integrates static code analysis with an AI-driven planner to generate precise migration instructions. A coordinator and encoder work in tandem, leveraging AI-generated, example-driven migration guides to achieve high-fidelity translation without requiring test code. Innovatively, an AI-based evaluator assesses migration quality, establishing a self-reinforcing development loop. Evaluated in real-world, large-scale production environments, the method accelerates framework migration by 6.4–8×, substantially expediting model infrastructure evolution.

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.

Problem

• model migration
• deep learning frameworks
• code maintenance
• TensorFlow to JAX
• AI-based products

Innovation

• multi-agent AI system
• model migration
• TensorFlow to JAX
• AI planner
• code transformation