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.
Developer documentation
API reference
- fedjax core
- Subpackages
- Federated algorithm
- Federated data
- Client dataset
- For each client
- Model
FederatedAlgorithmFederatedDataSubsetFederatedDataSQLiteFederatedDataInMemoryFederatedDataFederatedDataBuilderSQLiteFederatedDataBuilderClientPreprocessorshuffle_repeat_batch_federated_data()padded_batch_federated_data()RepeatableIteratorClientDatasetBatchPreprocessorbuffered_shuffle_batch_client_datasets()padded_batch_client_datasets()for_each_client()for_each_client_backend()set_for_each_client_backend()Modelcreate_model_from_haiku()create_model_from_stax()evaluate_model()model_grad()model_per_example_loss()evaluate_average_loss()ModelEvaluatorAverageLossEvaluatorgrad()
- fedjax.aggregators
- fedjax.algorithms
- fedjax.datasets
- fedjax.models
- fedjax.training