Google系 PyTorch TPU バックエンドが公開したもの
Google系の `google-pytorch/torch_tpu` は、PyTorch を TPU で動かすためのバックエンドを公開しているリポジトリだ。単に「コードがあります」という話ではなく、TPU 上で PyTorch のモデルや演算を回すための実装、ビルド手順、開発環境の整え方までひとまとめにしているのが特徴だ。しかも README の情報だけでも、実運用に近いところまでかなり踏み込んでいることが分かる。研究用の試作ではなく、実際に使う人をかなり意識した公開物だと思う。 このリポジトリは、Google の Tensor Processing Units、つまり TPU 向けに PyTorch のバックエンドを提供するプロジェクトとして説明されている。README では、PyTorch のモデルや演算を TPU デバイス上で実行できるようにすることが目的だと明記している。中核になるのは TPU 向けの ATen kernel 実装で、PyTorch の内部演算に相当する部分を TPU 側に持っていく構成だ。あわせて、XLA 計算の実行を最適化するための compilation
papoo.work