A Multi-agent AI System for Deep Learning Model Migration from TensorFlow to JAX
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2026
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866911550379393024 |
|---|---|
| author | Nikolov, Stoyan Konrad, Bernhard Gronbach, Moritz Kumar, Niket Yan, Ann Singh, Varun Liang, Yaning Ranganathan, Parthasarathy |
| author_facet | Nikolov, Stoyan Konrad, Bernhard Gronbach, Moritz Kumar, Niket Yan, Ann Singh, Varun Liang, Yaning Ranganathan, Parthasarathy |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2603_27296 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | A Multi-agent AI System for Deep Learning Model Migration from TensorFlow to JAX Nikolov, Stoyan Konrad, Bernhard Gronbach, Moritz Kumar, Niket Yan, Ann Singh, Varun Liang, Yaning Ranganathan, Parthasarathy Software Engineering Artificial Intelligence 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. |
| title | A Multi-agent AI System for Deep Learning Model Migration from TensorFlow to JAX |
| topic | Software Engineering Artificial Intelligence |
| url | https://arxiv.org/abs/2603.27296 |