AI支援コード移行：GoogleがTensorFlowからJAXへの移行を6倍高速化

Googleは、AI・インフラストラクチャチームが開発した革新的なマルチエージェントAIシステムにより、TensorFlowモデルからJAXへの移行速度を6倍以上向上させたと発表しました。この成果は、Google Cloud Next基調講演でSundar Pichai氏によって紹介されました。

AIコーディングエージェントがソフトウェア業界で急速に普及する中、Googleはこれらのツールを大規模なシステム全体のコード移行に活用しています。x86からARM、int32からint64、JUnit3からJUnit4、Joda-Timeからjava.timeへの移行で成果を挙げてきましたが、AIモデルの移行はこれまで以上に複雑です。TensorFlowからJAXへの移行は、単なる構文の更新ではなく、数千行のコードの整理、複数ファイルにわたる状態管理、厳密な数学的等価性の維持を必要とする長期的なタスクです。

この課題に対処するため、Googleチームは特殊化されたマルチエージェントアーキテクチャを開発しました：

• プランナーエージェント：コンパイラベースの静的解析を用いてコードベース全体の依存関係ツリーをマッピングし、リーフノード（未移行の依存関係がないレイヤー）から上方向への段階的な移行計画を策定します。
• オーケストレーターエージェント：プロジェクトマネージャーとして機能し、計画ステップを動的にグループ化してコンテキストウィンドウを最適化し、必要なドメイン知識を注入し、ビルド失敗時の復旧を処理します。
• コーダーエージェント：推論と行動を行うエージェントで、内部IDEツールに統合され、ファイルの読み取り、コードの書き込み、ビルドの実行、単体テストの実行が可能です。「テストと修正」のループで動作し、ターゲット言語でコンパイル可能で検証可能なコンポーネントを生成するまで自己修正を続けます。

移行品質を確保するため、チームはスケーラブルな階層型Playbookシステムを設計しました。一般的なリポジトリ指示から特定のクライアント（YouTubeのランキングモデルなど）向けの「黄金の例」までをカバーし、AIの幻覚を防止します。検証プロセスには、数学的検証（元のTFレイヤーと新しいJAXレイヤー間の最大誤差を勾配上昇法で見つける）と定性的評価（ブラインド監査LLM Judgeによるアーキテクチャチェックリストの評価）が含まれます。

実際のテストでは、このシステムはYouTubeの複雑なモデル（数千行のコード、数百のレイヤー、深いメトリクス依存関係を含む）に対して6.4倍から8倍の速度向上を達成しました。従来は数ヶ月かかっていた手動移行が、今では数週間のAI支援コード生成と専門家によるレビューに短縮されました。

Googleは、このアプローチがTensorFlowからJAXへの移行に限らず、任意のプログラミング言語やフレームワーク間の移行に適用可能なフレームワーク非依存の設計であると述べています。チームはプロセスの詳細を技術論文として公開しており、決定論的静的解析、厳格なテストループ、特殊化されたマルチエージェントアーキテクチャを組み合わせることで、業界で最も複雑なソフトウェアエンジニアリングの課題を安全に自動化できると強調しています。