PyTorch/XLA 2.4 предлагает значительные улучшения для глубокого обучения на ТПУ и GPU. В этом выпуске были внесены улучшения в Pallas, язык настроенных ядер, который теперь поддерживает и ТПУ, и GPU, что улучшает производительность с помощью кода на Python. Новые вызовы API, такие как torch_xla.sync(), упрощают интеграцию в существующие потоки PyTorch. Экспериментальный режим immediate позволяет немедленно выполнять операции на целевом оборудовании, хотя для ТПУ требуется вызов "mark_step" для эмуляции. Улучшения в Pallas включают поддержку Flash Attention и Paged Attention, а также встроенные ядра Megablocks для блочной разреженной матрицы умножения. Кроме того, новый интерфейс командной строки ТПУ, tpu-info, облегчает отладку, отображая информацию о загрузке и устройстве, аналогично инструменту Nvidia nvidia-smi. Несмотря на эти изменения, существующий код остается совместимым, что делает обновление незаметным для разработчиков. Эти обновления в целом направлены на улучшение удобства и производительности в PyTorch/XLA.