GoogleとIntelが共同開発、Intel GPUでJAXモデルを高速化するPJRT API登場：AIフレームワークへのハードウェア最適化が容易に

GoogleはPJRTの初のプラグインPJRT APIをIntelと共同開発したことを発表した。PJRT APIを実装することでIntel GPU上でのJAXモデルの実行や、XLAアクセラレーションを使用したTensorFlowおよびPyTorchモデルの初期Intel GPUサポートも可能になる。

[＠IT]

　Googleは2023年6月1日（米国時間）、Intel GPU上でJAXモデルをシームレスに実行する「Intel Extension for TensorFlow」で、初のPJRT（Performance JIT Runtime）プラグインの実装を発表した。PJRT APIは別途開発されたIntel GPUプラグインとの統合を簡素化し、JAXへの迅速な統合を可能にする。また、PJRTの実装により、XLAアクセラレーションを備えたTensorFlowとPyTorchモデルの初期Intel GPUサポートも有効になる。

Intel Data Center GPU Maxシリーズ（提供：Google）

　IntelとGoogleは共同でTensorFlow PluggableDeviceメカニズムを開発した。このメカニズムはTensorFlowを新しいデバイスに拡張するためにサポートされた手法であり、ハードウェアベンダーはこれを使用して、独自のプラグインバイナリのリリースが可能になる。Intelは、XLAコンパイラ用のモジュラーインタフェースの構築や、IntelGPU上でJAXワークロードを実行するためのPJRTプラグインの開発でGoogleと協力を続けている。

JAX

　JAXは、GPUやTPUなどのハイパフォーマンスコンピューティングデバイスでの複雑な数値計算のために設計されたオープンソースのPythonライブラリである。NumPy関数をサポートし、自動微分や、ニューラルネットワークを構築、訓練するためのコンポーザブル関数変換システムを提供する。