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

• URL: http://arxiv.org/abs/2603.27296v1
• Date: Sat, 28 Mar 2026 14:59:40 GMT
• Title: A Multi-agent AI System for Deep Learning Model Migration from TensorFlow to JAX
• Abstract summary: Google has been pioneering machine learning usage across dozens of products.<n>Maintaining the multitude of model source codes in different ML frameworks and versions is a significant challenge.<n>We describe an AI-based multi-agent system that we built to support automatic migration of deep learning models into JAX-based ones.
• Score: 1.4085663724094575
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• AIGen: Automating AI Bill of Materials Generation Through Hybrid MLOps Integration [4.706463542284726]
  This paper presents AIGen, a modular AIBoM generator that produces machine-readable AI system inventories.<n>A plugin interface allows practitioners to extend the tool with domain-specific collectors without modifying the core dependencies.
  arXiv  Detail & Related papers  (2026-07-29T09:13:00Z)
• Structuring agentic AI for HPC code modernization [0.0]
  NMAP-RKPM is a 60,000-line, 3D explicit solid mechanics physics engine based on the Reproducing Kernel Particle Method (RKPM)<n>We converted this single-threaded, Fortran based MPI application into a OpenMP-parallel C++ based MPI tool in the span of a few months.<n>The paper provides both the AI-assisted steps that were successful and the problems that we had to overcome, alongside the reasoning behind the chosen path.
  arXiv  Detail & Related papers  (2026-06-07T16:07:29Z)
• MLS-Bench: A Holistic and Rigorous Assessment of AI Systems on Building Better AI [92.82821272460423]
  MLS-Bench is a benchmark for evaluating whether AI systems can invent generalizable and scalable ML methods.<n>We study the effects of test-time scaling, adaptive compute allocation, and context provision on agents' discovery performance.
  arXiv  Detail & Related papers  (2026-05-09T04:29:46Z)
• AIBuildAI: An AI Agent for Automatically Building AI Models [31.05147308924019]
  We introduce AIBuildAI, an AI agent that automatically builds AI models from a task description and training data.<n>We evaluate AIBuildAI on MLE-Bench, a benchmark of realistic Kaggle-style AI development tasks.
  arXiv  Detail & Related papers  (2026-04-15T22:17:05Z)
• Agentic Artificial Intelligence (AI): Architectures, Taxonomies, and Evaluation of Large Language Model Agents [14.448267395835721]
  We propose a unified taxonomy that breaks agents into Perception, Brain, Planning, Action, Tool Use, and Collaboration.<n>We also group the environments in which these agents operate, including digital operating systems, embodied robotics, and other specialized domains.
  arXiv  Detail & Related papers  (2026-01-18T19:51:16Z)
• SysMoBench: Evaluating AI on Formally Modeling Complex Real-World Systems [12.181911851729614]
  We present SysMoBench, a benchmark that evaluates AI's ability to formally model large, complex systems.<n>We focus on concurrent and distributed systems, which are keystones of today's critical computing infrastructures.
  arXiv  Detail & Related papers  (2025-09-27T05:24:54Z)
• Continual Learning for Generative AI: From LLMs to MLLMs and Beyond [56.29231194002407]
  We present a comprehensive survey of continual learning methods for mainstream generative AI models.<n>We categorize these approaches into three paradigms: architecture-based, regularization-based, and replay-based.<n>We analyze continual learning setups for different generative models, including training objectives, benchmarks, and core backbones.
  arXiv  Detail & Related papers  (2025-06-16T02:27:25Z)
• Vibe Coding vs. Agentic Coding: Fundamentals and Practical Implications of Agentic AI [0.36868085124383626]
  Review presents a comprehensive analysis of two emerging paradigms in AI-assisted software development: vibe coding and agentic coding.<n> Vibe coding emphasizes intuitive, human-in-the-loop interaction through prompt-based, conversational interaction.<n>Agentic coding enables autonomous software development through goal-driven agents capable of planning, executing, testing, and iterating tasks with minimal human intervention.
  arXiv  Detail & Related papers  (2025-05-26T03:00:21Z)
• R&D-Agent: An LLM-Agent Framework Towards Autonomous Data Science [70.1638335489284]
  High-level machine learning engineering tasks remain labor-intensive and iterative.<n>We introduce R&D-Agent, a comprehensive, decoupled, and framework that formalizes the machine learning process.<n>R&D-Agent defines the MLE into two phases and six components, turning agent design for MLE into a principled, testable process.
  arXiv  Detail & Related papers  (2025-05-20T06:07:00Z)
• Generating a Low-code Complete Workflow via Task Decomposition and RAG [0.040964539027092926]
  GenAI-based systems are more difficult to design due to their scale and versatility.<n>We formalize two techniques, Task Decomposition and Retrieval-Augmented Generation, as design patterns for GenAI-based systems.<n>As these two patterns affect the entire AI development cycle, we explain how they impacted the dataset creation, model training, model evaluation, and deployment phases.
  arXiv  Detail & Related papers  (2024-11-29T20:13:56Z)
• Inference Optimization of Foundation Models on AI Accelerators [68.24450520773688]
  Powerful foundation models, including large language models (LLMs), with Transformer architectures have ushered in a new era of Generative AI. As the number of model parameters reaches to hundreds of billions, their deployment incurs prohibitive inference costs and high latency in real-world scenarios. This tutorial offers a comprehensive discussion on complementary inference optimization techniques using AI accelerators.
  arXiv  Detail & Related papers  (2024-07-12T09:24:34Z)
• Symbolic Learning Enables Self-Evolving Agents [55.625275970720374]
  We introduce agent symbolic learning, a systematic framework that enables language agents to optimize themselves on their own. Agent symbolic learning is designed to optimize the symbolic network within language agents by mimicking two fundamental algorithms in connectionist learning. We conduct proof-of-concept experiments on both standard benchmarks and complex real-world tasks.
  arXiv  Detail & Related papers  (2024-06-26T17:59:18Z)
• Automated Machine Learning: A Case Study on Non-Intrusive Appliance Load Monitoring [81.06807079998117]
  We propose a novel approach to enable Automated Machine Learning (AutoML) for Non-Intrusive Appliance Load Monitoring (NIALM)<n>NIALM offers a cost-effective alternative to smart meters for measuring the energy consumption of electric devices and appliances.
  arXiv  Detail & Related papers  (2022-03-06T10:12:56Z)
• Technology Readiness Levels for AI & ML [79.22051549519989]
  Development of machine learning systems can be executed easily with modern tools, but the process is typically rushed and means-to-an-end. Engineering systems follow well-defined processes and testing standards to streamline development for high-quality, reliable results. We propose a proven systems engineering approach for machine learning development and deployment.
  arXiv  Detail & Related papers  (2020-06-21T17:14:34Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.