FAX: Scalable and Differentiable Federated Primitives in JAX
Keith Rush Zachary Charles Zachary Garrett
We present FAX, a JAX-based library designed to support large-scale
distributed and federated computations in both data center and cross-device
applications. FAX leverages JAX's sharding mechanisms to enable native
targeting of TPUs and state-of-the-art JAX runtimes, including Pathways. FAX
embeds building blocks for federated computations as primitives in JAX. This
enables three key benefits. First, FAX computations can be translated to XLA
HLO. Second, FAX provides a full implementation of federated automatic
differentiation, greatly simplifying the expression of federated computations.
Last, FAX computations can be interpreted out to existing production
cross-device federated compute systems. We show that FAX provides an easily
programmable, performant, and scalable framework for federated computations in
the data center. FAX is available at
https://github.com/google-research/google-research/tree/master/fax .