One thing I really love about XLA is GSPMD which effectively allows scalable distributed training in practice. However, I was quite curious how it is related to matrix multiplication though, given XLA is more focusing on graph-level optimization and basically offloads matmul to other libraries like Triton and cuBLAS