AI辅助代码迁移：谷歌将TensorFlow迁移到JAX的速度提升了6倍

谷歌近日宣布，其AI和基础设施团队开发了一种创新的多智能体AI系统，可将TensorFlow模型迁移到JAX的速度提高6倍以上。这一成果在Google Cloud Next主题演讲中由Sundar Pichai重点介绍。

随着AI编码代理在软件行业的普及，谷歌正将这些工具应用于大规模系统级代码迁移。从x86到ARM、从int32到int64、从JUnit3到JUnit4等迁移工作已取得成效，但AI模型迁移的复杂性远超此前任何项目。TensorFlow到JAX的迁移不仅仅是语法转换，而是需要理清数千行代码、管理跨文件的复杂状态、保持数学等价性的长期任务。

为应对这一挑战，谷歌团队开发了专门的多智能体架构：

• 规划器代理：通过基于编译器的静态分析，映射代码库的完整依赖树，制定从叶节点向上的分步迁移计划。
• 协调器代理：充当项目经理，动态分组计划步骤，注入领域知识，处理构建失败时的恢复。
• 编码器代理：集成到内部IDE，具备读写文件、运行构建和单元测试的能力，采用“测试与修复”循环，直到产生可编译的组件。

为确保迁移质量，团队设计了可扩展的分层Playbook系统，从通用仓库指南到特定客户（如YouTube排名模型）的“黄金示例”，避免AI幻觉。验证环节包括：数学验证（通过梯度上升寻找原始层与新层之间的最大误差）和定性评估（使用盲审LLM Judge对迁移代码进行架构检查）。

在实际测试中，该系统对YouTube的复杂模型（包含数千行代码、数百层和深度指标依赖）实现了6.4至8倍的加速。原本需要数月的人工迁移工作，现在仅需几周AI辅助代码生成加上专家人工审查。

谷歌表示，这一方案不仅适用于TensorFlow到JAX的迁移，其框架无关的特性可扩展到任意编程语言或框架之间的迁移。团队已发表技术论文详细描述该流程，并强调结合确定性静态分析、严格测试循环和专门化多智能体架构，可以安全地自动化业界最复杂的软件工程挑战。