TPU/GPU collective primitives

Training LLMs on TPUs/GPUs is often constrained not by the speed of the matrix multiplication units (MXUs), but either by data transfer from high bandwidth memory (HBM) to MXU, or by data transfer between accelerators. A large part of optimizing training/inference programs consists of finding ways to overlap data transfer and matrix multiplication operations - transferring data that will be needed…