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

Computer Science > Software Engineering arXiv:2603.27296 (cs) [Submitted on 28 Mar 2026] Title: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 View a PDF of the paper titled A Multi-agent AI System for Deep Learning Model Migration from TensorFlow to JAX, by Stoyan Nikolov and 7 other authors View PDF HTML (experimental) 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 adh...