ALX: Large Scale Matrix Factorization on TPUs

Harsh Mehta, Steffen Rendle, Walid Krichene +1

We present ALX, an open-source library for distributed matrix factorization using Alternating Least Squares, written in JAX. Our design allows for efficient use of the TPU architec…