xla

收录理由

PyTorch/XLA 是让 PyTorch 模型在 Google Cloud TPU 上跑起来的主要通道。它把张量操作交给 XLA 编译器处理，编译器会对计算图做融合与优化，让 TPU 的硬件能力真正发挥出来，而不是沦为慢吞吞的兼容层。这个包可以直接用 pip 装到 TPU 虚拟机上，通过 PJRT 运行时支持跨芯片的分布式训练，像 Stable Diffusion 这种吃算力的负载也扛得住。很多场景下 TPU 每 flop 的成本比 GPU 低不少，对想保留现有 PyTorch 代码、又打算迁到 TPU 基础设施的团队来说，这是一条很实在的路。