AI を活用したコード移行の先駆け: Google が TensorFlow から JAX への移行を 6 倍高速化した方法

※この投稿は米国時間 2026 年 5 月 7 日に、Google Cloud blog に投稿されたものの抄訳です。

AI コーディング エージェントは、ソフトウェア業界全体で急速に普及しつつあり、デベロッパーが日常的にコードを記述、テスト、デバッグする方法を根本的に変えています。このツールは、ローカルで完結するタスクに優れた機能性を発揮しますが、大規模で体系的なコードベースの移行に適用するには、まったく新しいアプローチを必要とします。

Google には、多くの移行ワークフローに AI を組み込み、このような課題に対処してきた実績があります。たとえば、x86 から ARM への移行では Google Axion プロセッサでワークロードを有効にし、int32（32 ビット整数）識別子から int64（64 ビット整数）識別子への移行によって ID の枯渇を防ぎ、JUnit3 から JUnit4 への移行でテストを簡素化し、Joda-Time から java.time への移行で最新の時刻ライブラリを使用可能にしました。しかし、AI モデルの移行には新次元の複雑さがともない、AI による移行支援のためのさらに高度な手法が必要となります。

本番環境グレードの ML モデルを 1 つのフレームワークから別のフレームワークに変換する（たとえば、TensorFlow（TF）から JAX に変換する）ことは、単純な構文の更新ではありません。数千行のコードを解きほぐし、複数ファイルにまたがる複雑な状態管理を実現し、正確な数学的等価性を維持しなければならない、長期的な展望を必要とするタスクです。これは、汎用の単一エージェントによるコーディング アシスタントには過大な重圧となります。長いワークフローでは頻繁にコンテキストが失われ、API を巻き込んだハルシネーションが起きたり、リポジトリ全体でビルド可能なコード生成が行われなくなったりします。

この業界全体の問題に対し、Google AI / インフラストラクチャ チームは、世界に先駆けて新しいアプローチを開発しました。このアプローチは、モデルの移行を 6 倍速くするという、先日の Google Cloud Next の基調講演で Sundar が強調したマイルストーンを達成するものとなりました。この投稿では、Google がどのように専用マルチエージェント AI システムをデプロイして、Google 最大規模の本番環境モデルを TF から JAX に移行したのかについて、ご紹介します。

TF から JAX への移行を加速

Google の多くのチームでは（業界全体にも言えることですが）、未来に向けたスケーラブルな ML は JAX を基盤として構築されています。関数型でステートレスなパラダイムを中心に設計された JAX は、最新の Tensor Processing Unit（TPU）インフラストラクチャと XLA コンパイル向けに最適化されており、最新 AI スタックの強固な基盤となるものです。

これは未来への進歩であると同時に、非常に大きな課題も提示しています。TensorFlow は、オブジェクト指向でステートフル レイヤの初期化と静的実行グラフを特徴とするフレームワークであり、現在、何千もの本番環境モデルが、この TF を基盤に構築されています。これらのモデルを JAX に手動で移行するには、レイヤの相互作用や状態管理の明示的な方法を根本的に見直す必要があります。このような移行は、それだけで大規模な組織全体における数百（場合によっては数千）のソフトウェア エンジニアリング（SWE）年数を必要とします。組織にとっては、この時間を新しいアーキテクチャの調査やプロダクト イノベーションの推進に費やす方がはるかに有益です。

この課題を AI で克服する取り組みは、Google AI / インフラストラクチャ チームの野心的な実験として始まりましたが、今では社内全体の複雑なエンジニアリングの問題に対処可能な再現性のあるブループリントにまで進化しています。

単一エージェント コーディングからの脱却

エージェントによるコード変換は、初期の実験では単純なモデルへの有効性が確認されました。しかし、現実的な Google 規模の移行に直面したとき、つまり、複数のファイルにまたがる数千行のコードの複雑な本番環境グレードのモデルを扱うとき、汎用的な単一エージェントでは困難を極めました。高次元の構造ルールと現実のこまごまとした実行上の問題の折り合いがつかず、重要なファイルを上書きしたり、必要な機能をとばしたりと、さまざまな障害が発生しました。このような大規模な移行によくある課題を克服するために、Google は以下で構成される高度に専門化されたマルチエージェント アーキテクチャを開発しました。

• Planner エージェント: 決定論的なコンパイラベースの静的分析を使用して、コードベースの依存関係全体をツリー構造にマッピングします。その後、他のエージェントと連携して移行を個別の段階的プランに分割し、「リーフノード」（依存関係のない未移行レイヤ）から上に向かって論理的に移行が行われるようにします。

• Orchestrator エージェント: このエージェントはプロジェクト マネージャーとして機能します。プラン ステップを管理しやすいチャンクに動的にグループ化してコンテキスト ウィンドウの対象を絞り、必要なドメイン知識を注入します。ステップがビルドされない場合は障害復旧を適用します。

• Coder エージェント: 推論と行動の役割を担う主力エージェントです。社内の IDE ツールに直接統合され、ファイルの読み取り、コードの記述、ビルドの実行、単体テストの実施が可能です。「テストと修正」のループで動作し、コンパイル可能で検証可能なコンポーネントをターゲット言語で生成するまで自己修正を繰り返せる点が最大の特長です。

スケーラブルな検証と柔軟な Playbook

生成 AI モデルの性能は、提供されるコンテキストの質に左右されます。移行元と移行先のアーキテクチャが 1 対 1 でマッピングされることはほとんどないため、Google はスケーラブルで階層的な一連の Playbook を策定しました。

これらの Playbook は、一般的なリポジトリの手順から、手動移行の成功事例から抽出された非常に具体的な「模範例」まで多岐にわたります。クライアント固有の Playbook（たとえば、YouTube 独自のランキング モデル インフラストラクチャ向けに調整されたもの）を Orchestrator にフィードすることで、一般的なハルシネーションが回避され、組織内のコーディング基準が厳守されます。この Playbook の構成は特定のフレームワークに依存しないため、任意の 2 つのプログラミング言語やフレームワーク間の移行をガイドするように調整することができます。

さらに、生成されたコードが実際にプロダクション レディであることを確認するために、以下の厳格な品質指標も設定されています。

1. 定量的検証: コードの正確性をユニットごとに数学的に検証します。TF から JAX への移行の場合、勾配上昇法のアルゴリズムを使用して元の TF レイヤと新しい JAX レイヤの間の最大誤差を検出し、関数的な同等性を数学的に検証します。

2. 定性的評価: 移行されたコードを、一連の定性的基準に照らして評価します。TF から JAX への移行の場合、盲検監査の LLM Judge をデプロイして、移行後のコードをフレームワークに依存しないアーキテクチャ チェックリストに照らして採点します。これにより、ドメイン固有の重要なロジックを確実に把握することができます。

移行のスピードを革新

このマルチエージェント システムをデプロイすることで、ソフトウェア移行の経済的効率性は劇的に向上します。

実際の複雑な YouTube モデル（数千行のコード、数百のレイヤ、複雑なメトリック依存関係を含む）で評価したところ、マルチエージェント システムは手動での移行よりも 6.4 倍から 8 倍の高速化を実現しました。従来は数か月分のソフトウェア エンジニアリング（SWE）作業が必要だったものが、わずか数週間の AI によるコード生成と、その後の専門家（人間）によるレビューで完了できるようになったのです。

ボイラープレートの効果的な処理、ターゲット言語のイディオム特定、依存関係のマッピング、単体テストの生成が自動的に行われ、エンジニアは手動のコード変換担当者ではなく、レビュアーやアーキテクトとしての役割を担えるようになります。

AI を活用する時代の展望

AI は技術革新のペースを変革しています。大規模な移行を推進する力は AI で加速する必要があります。組織はそれなしでは、手動の作業に追われ、最新のブレークスルーを取り入れたり、システムのセキュリティ、信頼性、パフォーマンスを維持したりすることが難しく、そのギャップはますます広がっていきます。

ML 実装を 1 つの ML フレームワークから別の ML フレームワークに移行する Google の取り組みは、決定論的静的分析、厳格なテストループ、特殊なマルチエージェント アーキテクチャを組み合わせることで、業界でも特に複雑なソフトウェア エンジニアリングの課題を安全に自動化できることを示しています。プロセスの詳細については、こちらの技術論文をご覧ください。

^{この取り組みは、Google 全体でのコラボレーションの成果です。主な貢献者である Stoyan Nikolov、Niyati Parameswaran、Bernhard Konrad、Moritz Gronbach、Niket Kumar、Ann Yan、Varun Singh、Yaning Liang、Antoine Baudoux、Xevi Miró Bruix、Daniele Codecasa、Madhura Dudhgaonkar、Elian Dumitru、Alex Ivanov、Christopher Milne-O’Grady、Ahmed Omran、Ivan Petrychenko、Assaf Raman、Stefan Schnabl、Yurun Shen、Maxim Tabachnyk、Niranjan Tulpule、Amin Vahdat、Jeff Zhou に感謝いたします。}

- ドメイン応用 ML / AI / インフラストラクチャ プロダクト責任者、Jamie Rogers

- Google Fellow 兼 AI およびインフラストラクチャ担当バイス プレジデント、Parthasarathy Ranganathan