JAXMg: A multi-GPU linear solver in JAX

Jacob Tutt, Roeland Wiersema

Solving large dense linear systems and eigenvalue problems is a core requirement in many areas of scientific computing, but scaling these operations beyond a single GPU remains cha…