Google open-sources TorchTPU, a native PyTorch back end for TPUs
Проект TorchTPU предоставляет бэкенд PyTorch для тензорных процессоров (TPU) Google, позволяя запускать модели и операции PyTorch напрямую на TPU. Он включает собственные ядра ATen для различных операций, кэш компиляции для оптимизации вычислений XLA, поддержку распределённого обучения (параллелизм по данным и тензорам) и интеграцию с torch.compile() для графового режима. В репозитории есть примеры для Llama, Qwen, ResNet, DLRM, minGPT и UNet, а также утилиты для отладки и бенчмаркинга.
Установить последнюю версию (0.2.0) можно через pip с дополнительным индексом PyTorch CPU; поддерживается Python 3.11+. Сборка из исходников ведётся через Bazel: обязателен флаг --config=wheel_common, иначе колесо упадёт при импорте из-за дублирующихся статических регистраций. Без учётных данных удалённого исполнения (RBE) нужно добавить --config=no_rbe строго после wheel_common. Для разработки на C++ предусмотрен скрипт setup_clangd.py, генерирующий compile_commands.json (первый запуск может занять более 30 минут). Проект распространяется под лицензией Apache 2.0 и открыт для вкладов.