title: FAX: Scalable and Differentiable Federated Primitives in JAX
summary: 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 .
paper_authors._id: 65f11795c8c8fba63192afc4
paper_authors.name: Keith Rush
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-03-13T16:15:00.532Z
paper_authors._id: 65f1cf3aea7a16c6295f4737
paper_authors.avatarUrl: /avatars/905106897916e87e114ccc2d33e8e1c2.svg
paper_authors.fullname: Keith Rush
paper_authors.type: user
paper_authors.user: jkr26
paper_authors._id: 65f11795c8c8fba63192afc5
paper_authors.name: Zachary Charles
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-03-14T08:09:56.031Z
paper_authors._id: 65f1c0bf22ea2a37a28b6466
paper_authors.avatarUrl: /avatars/d9a11158bad338ada0d044476d8681c1.svg
paper_authors.fullname: Zachary Charles
paper_authors.type: user
paper_authors.user: zcharles8
paper_authors._id: 65f11795c8c8fba63192afc6
paper_authors.name: Zachary Garrett
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-03-14T08:10:16.343Z
paper_authors._id: 65f1be9109a3517bd1c8e17b
paper_authors.avatarUrl: /avatars/bc4112a3f749fe969ebb633b2e0cbc14.svg
paper_authors.fullname: Zachary Garrett
paper_authors.type: user
paper_authors.user: ZacharyGarrett
paper_title: FAX: Scalable and Differentiable Federated Primitives in JAX
paper_summary: 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 .
paper_ai_keywords: JAX-based library large-scale distributed and federated computations TPUs Pathways sharding mechanisms federated computations XLA HLO federated automatic differentiation cross-device federated compute systems easily programmable framework