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的遷移，其框架無關的特性可擴展到任意編程語言或框架之間的遷移。團隊已發表技術論文詳細描述該流程，並強調結合確定性靜態分析、嚴格測試循環和專門化多智能體架構，可以安全地自動化業界最複雜的軟件工程挑戰。