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

arXiv: https://arxiv.org/abs/2603.27296 · HTML: https://mlsanity.com/paper/2603.27296/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 10 of its 24 references that are in the map.
- First posted on arXiv 28 Mar 2026.

## Builds on (most cited first, 10 of 10)

- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [Agentless: Demystifying LLM-based Software Engineering Agents](https://mlsanity.com/paper/2407.01489/) (2024)
- [Agentic Context Engineering: Evolving Contexts for Self-Improving Language Models](https://mlsanity.com/paper/2510.04618/) (2025)
- [Dynamic Cheatsheet: Test-Time Learning with Adaptive Memory](https://mlsanity.com/paper/2504.07952/) (2025)
- [HyperAgent: Generalist Software Engineering Agents to Solve Coding Tasks at Scale](https://mlsanity.com/paper/2409.16299/) (2024)
- [SWE-EVO: Benchmarking Coding Agents in Long-Horizon Software Evolution Scenarios](https://mlsanity.com/paper/2512.18470/) (2025)
- [MaintainCoder: Maintainable Code Generation Under Dynamic Requirements](https://mlsanity.com/paper/2503.24260/) (2025)
- [RefAgent: A Multi-agent LLM-based Framework for Automatic Software Refactoring](https://mlsanity.com/paper/2511.03153/) (2025)
- [How is Google using AI for internal code migrations?](https://mlsanity.com/paper/2501.06972/) (2025)
- [FreshBrew: A Benchmark for Evaluating AI Agents on Java Code Migration](https://mlsanity.com/paper/2510.04852/) (2025)

## Built on by (most cited first, 2 of 2)

- [Evaluating Language Models on Cross-Language Code Functional Equivalence](https://mlsanity.com/paper/2608.23961/) (2026)
- [Agentic Framework for Deep Learning workload migration via In-Context Learning](https://mlsanity.com/paper/2606.15994/) (2026)

## Abstract (arXiv)

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.