FedJAX documentation

FedJAX is a library for developing custom Federated Learning (FL) algorithms in JAX. FedJAX prioritizes ease-of-use and is intended to be useful for anyone with knowledge of NumPy.

Getting started

• Federated datasets
• Working with models in FedJAX
• Federated learning algorithms

Developer documentation

• Contributing
• Building from source
• Running the tests
• Updating the docs

API reference

• fedjax core
  • Subpackages
    • fedjax.metrics
    • fedjax.optimizers
    • fedjax.tree_util
  • Federated algorithm
  • Federated data
  • Client dataset
  • For each client
  • Model
• fedjax.aggregators
  • Mean
  • Quantization
• fedjax.algorithms
  • AgnosticFedAvg
  • FedAvg
  • HypCluster
  • Mime
  • MimeLite
• fedjax.datasets
  • CIFAR-100
  • EMNIST
  • Shakespeare
  • Stack Overflow
• fedjax.models
  • EMNIST
  • Shakespeare
  • Stack Overflow
  • Toy Regression
• fedjax.training
  • Federated experiment
  • Tasks
  • Structured flags