Using Distributed Dense Matrix-Matrix Multiplication (matmul) in JAX

[kabicm]

As somebody who is coming from an HPC background, I am wondering if JAX would benefit from a fast, distributed, dense matrix-matrix multiplication algorithm like COSMA?

COSMA is GPU-accelerated and can also take advantage of fast gpu-to-gpu interconnects through NCCL or gpu-aware MPI. It is already used in some scientific applications in supercomputing centres.

I would be happy to integrate it into JAX if you think it would be beneficial?

Thanks!