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