Google открыла исходный код TorchTPU — нативного бэкенда PyTorch для TPU
Google open-sources TorchTPU, a native PyTorch back end for TPUs
Google выложила в открытый доступ TorchTPU — бэкенд, позволяющий запускать модели и операции PyTorch напрямую на тензорных процессорах TPU. Проект включает кастомные ядра ATen, кэш компиляции для оптимизации XLA-вычислений, поддержку распределённого обучения (data и tensor parallelism) и интеграцию с torch.compile(). В репозитории есть примеры для Llama, Qwen, ResNet, DLRM, minGPT и UNet, а также утилиты для отладки и бенчмаркинга. Установка доступна через pip (версия 0.2.0) или сборкой из исходников на Bazel.
Этот проект предоставляет бэкенд PyTorch для тензорных процессоров Google (TPU), позволяя пользователям запускать модели и операции PyTorch на устройствах TPU.