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 .

FAX是一个基于JAX的库，旨在支持数据中心和跨设备应用中的大规模分布式和联合计算。FAX利用JAX的分片机制，实现了针对TPUs和JAX运行时的本地化支持。FAX将联合计算的构建块嵌入JAX的原语中，并提供了对XLA HLO的翻译，联合自动微分功能的完整实现，以及对现有生产跨设备联合计算系统的解释。我们展示了FAX在数据中心中提供了易于编程、高性能和可扩展的联合计算框架。

FAX：具规模和可微分的JAX联邦原语