PyTorch/XLA 2.4 ofrece mejoras significativas para el aprendizaje profundo en TPU y GPU. Esta versión introduce mejoras en Pallas, un lenguaje de kernel personalizado, que ahora admite tanto TPU como GPU, lo que mejora el rendimiento con código Python. Las nuevas llamadas de API, como torch_xla.sync(), simplifican la integración en flujos de trabajo existentes de PyTorch. Un modo ávido experimental permite la ejecución inmediata de operaciones en hardware de destino, aunque los TPU requieren una llamada "mark_step" para la emulación. Las mejoras en Pallas incluyen soporte para Flash Attention y Paged Attention, y kernels de matrices esparsas de bloques integrados para la multiplicación de matrices de grupo. Además, una nueva interfaz de línea de comandos TPU, tpu-info, facilita la depuración al mostrar información de utilización y dispositivo, similar a la herramienta nvidia-smi de Nvidia. A pesar de estos cambios, el código existente permanece compatible, lo que hace que la actualización sea indolora para los desarrolladores. Estas actualizaciones buscan colectivamente mejorar la usabilidad y el rendimiento en PyTorch/XLA.