<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Federated datasets
This tutorial introduces datasets in FedJAX and how to work with them. By completing this tutorial, we'll learn about the best practices for working with datasets.
NOTE
Step1:  What are datasets in federated learning?
In the context of federated learning (FL), data is decentralized across clients, with each client having their own local set of examples. In light of this, we refer to two levels of organization for datasets
Step2:  fedjax.FederatedData provides methods for accessing metadata about the federated dataset, like the total number of clients, client ids, and number of examples for each client.
As seen in the output, there are 715 total clients in the Shakespeare dataset.
Each client has a unique client ID that can be used to query metadata about that client such as the number of examples that client has.
Step3:  As we notice, the client ids start with a random set of bits. This is to ensure that one can easily slice a FederatedData to obtain a random subset. While the client ids returned above look sorted, there is no such guarantee in general.
Step4:  Client datasets
We can query the dataset for a client from a federated dataset using their client ID and fedjax.FederatedData.get_client(). The output of fedjax.FederatedData.get_client() is a fedjax.ClientDataset. A fedjax.ClientDataset object
Stores all the examples for a given client and any preprocessing that should be applied on batches.
Provides methods for batching, shuffling, and iterating over preprocessed examples.
In other words, ClientDataset = examples + preprocessor.
Step5:  FedJAX assumes that an individual client dataset is small and can easily fit in memory. This assumption is also reflected in many of FedJAX's design decisions. The examples in a client dataset can be viewed as a table, where the rows are
the individual examples, and the columns are the features (labels are viewed as
a feature in this context).
FedJAX uses a column based representation when loading a dataset into memory.
Each column is a NumPy array x of rank at least 1, where x[i, ...] is
    the value of this feature for the i-th example.
The complete set of examples is a dict-like object, from str feature
    names, to the corresponding column values.
Traditionally, a row based representation is used for representing the entire
dataset, and a column based representation is used for a single batch.
In the context of federated learning, an individual client dataset is small
enough to easily fit into memory so the same representation is used for the
entire dataset and a batch.
The preprocessed examples of a single client can be viewed by calling all_examples().
Step6:  For Shakespeare, we are training a character-level language model, where the task is next character prediction, so the features are
Step7:  However, you probably noticed that x and y are arrays of integers not text. This is because  fedjax.datasets.shakespeare.load_data() does some minimal preprocessing, such as a simple character look up that mapped characters to integer IDs. Later, we'll go over how this preprocessing was applied and how to add your own custom preprocessing.
For comparison, here's the unprocessed version of the same client dataset.
Step8:  Accessing client datasets from fedjax.FederatedData
fedjax.FederatedData.get_client() works well for querying data for a single client for exploring the dataset. However, we often want to query for multiple client datasets at the same time. In most FL algorithms, tens to hundreds of clients particpate in each training round. If we are not careful, our code can spend a lot of time loading data, leaving the accelerators (GPU or TPU) to idle.
In light of this, fedjax.FederatedData offers more efficient methods for querying multiple client datasets.
We'll go through each access method from the MOST efficient to the LEAST efficient.
clients() and shuffled_clients()
Fastest sequential read friendly access. As we stated earlier, the client ids are appended with random bits. Hene, even sequential reads will go over clients in a pseudo-random order.
Step9:  They are generators, so we iterate over them to get the individual client datasets as tuples of (client_id, client_dataset).
clients() returns clients in an unspecified deterministic order. It is useful for going over the entire federated dataset for evaluation.
Step10:  shuffled_clients() provides a stream of infinitely repeating shuffled client datasets, using buffered shuffling. It is suitable for training rounds where a nearly random shuffling is good enough.
Step11:  get_clients()
Slower than clients() since it requires random read but uses prefetching to hide the latency of random read access. This also returns a generator of tuples of (client_id, client_dataset), in the order of the input client_ids.
Step12:  get_client()
Slowest way of accessing client datasets. We usually reserve this method only for interactive exploration of a small number of clients.
Step13:  Batching client datasets
Next we'll go over different methods of iterating over a fedjax.ClientDataset as batched examples. All the following methods can be invoked in 2 ways
Step14:  batch() for illustrations
Produces preprocessed batches in a fixed sequential order.
The final batch may contain fewer than batch_size examples. If used directly,
that may result in a large number of JIT recompilations. Therefore we
should use padded_batch() or shuffle_repeat_batch() instead in most scenarios.
Note here we are not talking about padding within an example, often done in processing sequence examples (e.g. the 0 labels below), but rather padding with "empty" examples in a batch.
Step15:  padded_batch() for evaluation
Produces preprocessed padded batches in a fixed sequential order for evaluation.
When the number of examples in the dataset is not a multiple of batch_size,
the final batch may be smaller than batch_size. This may lead to a large
number of JIT recompilations. This can be circumvented by padding the final
batch to a small number of fixed sizes controlled by num_batch_size_buckets.
Step16:  All batches contain an extra bool feature keyed by '__mask__'.
batch['__mask__'][i] tells us whether the i-th example in this batch
is an actual example (batch['__mask__'][i] == True), or a padding
example (batch['__mask__'][i] == False).
We repeatedly halve the batch size up to num_batch_size_buckets - 1 times, until
we find the smallest one that is also >= the size of the final batch. Therefore
if batch_size &lt; 2^num_batch_size_buckets, fewer bucket sizes will be actually
used. This will be seen when we look at the final batch that only has 4 examples when the original batch size was 8.
Step17:  shuffle_repeat_batch() for training
Produces preprocessed batches in a shuffled and repeated order for training.
Shuffling is done without replacement, therefore for a dataset of N examples,
the first ceil(N/batch_size) batches are guarranteed to cover the entire
dataset. Unlike batch() or padded_batch(), batches from
shuffle_repeat_batch() always contain exactly batch_size examples. Also
unlike TensorFlow, that holds even when drop_remainder=False.
By default the iteration stops after the first epoch.
Step18:  The number of batches produced from the iteration can be controlled by the (num_epochs, num_steps,
drop_remainder) combination
Step19:  If num_epochs is set and num_steps is None, as few batches as needed to go
over the dataset this many passes are produced. Further,
If drop_remainder is False (the default), the final batch is filled with
    additionally sampled examples to contain batch_size examples.
If drop_remainder is True, the final batch is dropped if it contains fewer
    than batch_size examples. This may result in examples being skipped when
    num_epochs=1.
Step20:  If num_steps is set and num_steps is None, exactly this many batches are
produced. drop_remainder has no effect in this case.
Step21:  If both num_epochs and num_steps are set, the fewer number of batches
between the two conditions are produced.
Step22:  If reproducible iteration order is desired, a fixed seed can be used. When
seed is None, repeated iteration over the same object may produce batches in a
different order.
Step23:  Preprocessing
Going from the raw features to features in a batch of examples ready for use in training/evalution often requires a few steps of preprocessing. Sometimes, we also want to add new preprocessing to an existing FederatedData.
Before going into the details of preprocessing, please note once again that all dataset related operations should be implemented in standard NumPy, not in JAX.
Preprocessing a batch of examples
The easiest way to add an additional preprocessing step is by appending a function that transforms a batch of examples to a FederatedData's list of preprocessing transformations on batches.
Below, we add a new z feature that stores the parity of y for our Shakespeare examples.
Step24:  Preprocessing at the client level
Sometimes we also need to do some preprocessing for the entire client dataset. For example, in the Shakespeare dataset, the raw features are just text strings, so we need to turn them into a NumPy array of chunks of integer labels just to be able to meaningfully batch them at all.
In most circumstances, adding a preprocessing step at the client level is unnecessary, and should be avoided, because the new preprocessing step added by preprocess_client is inserted into the middle of a chain of steps, just before all the batch level preprocessing steps registered by preprocess_batch. If not done carefully, a custom client level preprocessing can easily break the preprocessing chain.
Nevertheless, here's an example of client level processing for certain rare cases.
Step25:  The actual client level preprocessing in the Shakespeare dataset is a bit involved, so let's do something simpler
Step26:  Now, we can add another batch level preprocessor to produce x and y labels.
Step27:  In memory federated datasets
In many scenarios, it is desirable to create a small custom federated dataset from a collection of NumPy arrays for quick experimentation. FedJAX provides fedjax.InMemoryFederatedData to create small custom datasets. fedjax.InMemoryFederatedData takes a dictionary of numpy examples keyed by client id and creates a fedjax.FederatedData that is compatible with the rest of the library. We illustrate this below with a simple example.
<ASSISTANT_TASK:>
Python Code:
# Uncomment these to install fedjax.
# !pip install fedjax
# !pip install --upgrade git+https://github.com/google/fedjax.git
# !pip install tensorflow_datasets
import functools
import itertools
import fedjax
import numpy as np
Explanation: Federated datasets
This tutorial introduces datasets in FedJAX and how to work with them. By completing this tutorial, we'll learn about the best practices for working with datasets.
NOTE: For datasets, we operate over NumPy arrays NOT JAX arrays.
End of explanation
# We cap max sentence length to 8.
train_fd, test_fd = fedjax.datasets.shakespeare.load_data(sequence_length=8)
Explanation: What are datasets in federated learning?
In the context of federated learning (FL), data is decentralized across clients, with each client having their own local set of examples. In light of this, we refer to two levels of organization for datasets:
Federated dataset: A collection of clients, each with their own local datasets and metadata
Client dataset: The set of local examples for a particular client
We can think of federated data as a mapping from client ids to client datasets and client datasets as a list of examples.
federated_data = {
  'client0': ['a', 'b', 'c'],
  'client1': ['d', 'e'],
}
Federated datasets
FedJAX comes packaged with multiple federated datasets, and we will look at the Shakespeare dataset as an example. The Shakespeare dataset is created from The Complete Works of Shakespeare, by treating each character in the play as a "client", and their dialogue lines as the examples.
FedJAX organizes federated datasets as Python modules. load_data() from a dataset module loads all predefined splits together as fedjax.FederatedData objects, the interface for accessing all federated datasets. In the case of the Shakespeare dataset, load_data() returns two splits: train and test.
End of explanation
print('num_clients =', train_fd.num_clients())
# train_fd.client_ids() is a generator of client ids.
# itertools has efficient and convenient functions for working with generators.
for client_id in itertools.islice(train_fd.client_ids(), 3):
  print('client_id =', client_id)
  print('# examples =', train_fd.client_size(client_id))
Explanation: fedjax.FederatedData provides methods for accessing metadata about the federated dataset, like the total number of clients, client ids, and number of examples for each client.
As seen in the output, there are 715 total clients in the Shakespeare dataset.
Each client has a unique client ID that can be used to query metadata about that client such as the number of examples that client has.
End of explanation
# Slicing are based on the lexicographic order of client ids.
train_fd_0 = train_fd.slice(start=b'0', stop=b'1')
print('num_clients whose id starts with 0 =', train_fd_0.num_clients())
Explanation: As we notice, the client ids start with a random set of bits. This is to ensure that one can easily slice a FederatedData to obtain a random subset. While the client ids returned above look sorted, there is no such guarantee in general.
End of explanation
client_id = b'105f96df763d4ddf:ALL_S_WELL_THAT_ENDS_WELL_GUIDERIUS_AND_ARVIRAGUS'
client_dataset = train_fd.get_client(client_id)
print(client_dataset)
Explanation: Client datasets
We can query the dataset for a client from a federated dataset using their client ID and fedjax.FederatedData.get_client(). The output of fedjax.FederatedData.get_client() is a fedjax.ClientDataset. A fedjax.ClientDataset object
Stores all the examples for a given client and any preprocessing that should be applied on batches.
Provides methods for batching, shuffling, and iterating over preprocessed examples.
In other words, ClientDataset = examples + preprocessor.
End of explanation
client_dataset.all_examples()
Explanation: FedJAX assumes that an individual client dataset is small and can easily fit in memory. This assumption is also reflected in many of FedJAX's design decisions. The examples in a client dataset can be viewed as a table, where the rows are
the individual examples, and the columns are the features (labels are viewed as
a feature in this context).
FedJAX uses a column based representation when loading a dataset into memory.
Each column is a NumPy array x of rank at least 1, where x[i, ...] is
    the value of this feature for the i-th example.
The complete set of examples is a dict-like object, from str feature
    names, to the corresponding column values.
Traditionally, a row based representation is used for representing the entire
dataset, and a column based representation is used for a single batch.
In the context of federated learning, an individual client dataset is small
enough to easily fit into memory so the same representation is used for the
entire dataset and a batch.
The preprocessed examples of a single client can be viewed by calling all_examples().
End of explanation
examples = client_dataset.all_examples()
print('x', examples['x'][0])
print('y', examples['y'][0])
Explanation: For Shakespeare, we are training a character-level language model, where the task is next character prediction, so the features are:
x is a list of right-shifted sentences, e.g. sentence[:-1]
y is a list of left-shifted sentences, e.g. sentence[1:]
This way, y[i][j] corresponds to the next character after x[i][j].
End of explanation
# Unlike load_data(), load_split() always loads a single unprocessed split.
raw_fd = fedjax.datasets.shakespeare.load_split('train')
raw_fd.get_client(client_id).all_examples()
Explanation: However, you probably noticed that x and y are arrays of integers not text. This is because  fedjax.datasets.shakespeare.load_data() does some minimal preprocessing, such as a simple character look up that mapped characters to integer IDs. Later, we'll go over how this preprocessing was applied and how to add your own custom preprocessing.
For comparison, here's the unprocessed version of the same client dataset.
End of explanation
# clients() and shuffled_clients() are sequential read friendly.
clients = train_fd.clients()
shuffled_clients = train_fd.shuffled_clients(buffer_size=100, seed=0)
print('clients =', clients)
print('shuffled_clients =', shuffled_clients)
Explanation: Accessing client datasets from fedjax.FederatedData
fedjax.FederatedData.get_client() works well for querying data for a single client for exploring the dataset. However, we often want to query for multiple client datasets at the same time. In most FL algorithms, tens to hundreds of clients particpate in each training round. If we are not careful, our code can spend a lot of time loading data, leaving the accelerators (GPU or TPU) to idle.
In light of this, fedjax.FederatedData offers more efficient methods for querying multiple client datasets.
We'll go through each access method from the MOST efficient to the LEAST efficient.
clients() and shuffled_clients()
Fastest sequential read friendly access. As we stated earlier, the client ids are appended with random bits. Hene, even sequential reads will go over clients in a pseudo-random order.
End of explanation
# We use itertools.islice to select first three clients.
for client_id, client_dataset in itertools.islice(clients, 3):
  print('client_id =', client_id)
  print('# examples =', len(client_dataset))
Explanation: They are generators, so we iterate over them to get the individual client datasets as tuples of (client_id, client_dataset).
clients() returns clients in an unspecified deterministic order. It is useful for going over the entire federated dataset for evaluation.
End of explanation
print('shuffled_clients()')
for client_id, client_dataset in itertools.islice(shuffled_clients, 3):
  print('client_id =', client_id)
  print('# examples =', len(client_dataset))
Explanation: shuffled_clients() provides a stream of infinitely repeating shuffled client datasets, using buffered shuffling. It is suitable for training rounds where a nearly random shuffling is good enough.
End of explanation
client_ids = [
    b'1db830204507458e:THE_TAMING_OF_THE_SHREW_SEBASTIAN',
    b'140784b36d08efbc:PERICLES__PRINCE_OF_TYRE_GHOST_OF_VAUGHAN',
    b'105f96df763d4ddf:ALL_S_WELL_THAT_ENDS_WELL_GUIDERIUS_AND_ARVIRAGUS'
]
for client_id, client_dataset in train_fd.get_clients(client_ids):
  print('client_id =', client_id)
  print('# examples =', len(client_dataset))
Explanation: get_clients()
Slower than clients() since it requires random read but uses prefetching to hide the latency of random read access. This also returns a generator of tuples of (client_id, client_dataset), in the order of the input client_ids.
End of explanation
client_id = b'1db830204507458e:THE_TAMING_OF_THE_SHREW_SEBASTIAN'
print('client_id =', client_id)
print('# examples =', len(train_fd.get_client(client_id)))
Explanation: get_client()
Slowest way of accessing client datasets. We usually reserve this method only for interactive exploration of a small number of clients.
End of explanation
client_id = b'105f96df763d4ddf:ALL_S_WELL_THAT_ENDS_WELL_GUIDERIUS_AND_ARVIRAGUS'
client_dataset = train_fd.get_client(client_id)
Explanation: Batching client datasets
Next we'll go over different methods of iterating over a fedjax.ClientDataset as batched examples. All the following methods can be invoked in 2 ways:
Using a hyperparams object: This is the recommended way in library code. batch_fn(hparams).
Using keyword arguments: The keyword arguments are used to construct a new hyperparams object, or override an existing one. batch_fn(batch_size=2) or batch_fn(hparams, batch_size=2) to override batch_size.
For the most part, we'll use method 2 for this colab, but method 1 is more suitable for writing library code.
End of explanation
batches = list(client_dataset.batch(batch_size=8))
batches[-1]
Explanation: batch() for illustrations
Produces preprocessed batches in a fixed sequential order.
The final batch may contain fewer than batch_size examples. If used directly,
that may result in a large number of JIT recompilations. Therefore we
should use padded_batch() or shuffle_repeat_batch() instead in most scenarios.
Note here we are not talking about padding within an example, often done in processing sequence examples (e.g. the 0 labels below), but rather padding with "empty" examples in a batch.
End of explanation
# use list() to consume generator and store in memory.
padded_batches = list(client_dataset.padded_batch(batch_size=8, num_batch_size_buckets=3))
print('# batches =', len(padded_batches))
padded_batches[0]
Explanation: padded_batch() for evaluation
Produces preprocessed padded batches in a fixed sequential order for evaluation.
When the number of examples in the dataset is not a multiple of batch_size,
the final batch may be smaller than batch_size. This may lead to a large
number of JIT recompilations. This can be circumvented by padding the final
batch to a small number of fixed sizes controlled by num_batch_size_buckets.
End of explanation
padded_batches[-1]
Explanation: All batches contain an extra bool feature keyed by '__mask__'.
batch['__mask__'][i] tells us whether the i-th example in this batch
is an actual example (batch['__mask__'][i] == True), or a padding
example (batch['__mask__'][i] == False).
We repeatedly halve the batch size up to num_batch_size_buckets - 1 times, until
we find the smallest one that is also >= the size of the final batch. Therefore
if batch_size &lt; 2^num_batch_size_buckets, fewer bucket sizes will be actually
used. This will be seen when we look at the final batch that only has 4 examples when the original batch size was 8.
End of explanation
print('# batches')
len(list(client_dataset.shuffle_repeat_batch(batch_size=8)))
Explanation: shuffle_repeat_batch() for training
Produces preprocessed batches in a shuffled and repeated order for training.
Shuffling is done without replacement, therefore for a dataset of N examples,
the first ceil(N/batch_size) batches are guarranteed to cover the entire
dataset. Unlike batch() or padded_batch(), batches from
shuffle_repeat_batch() always contain exactly batch_size examples. Also
unlike TensorFlow, that holds even when drop_remainder=False.
By default the iteration stops after the first epoch.
End of explanation
infinite_bs = client_dataset.shuffle_repeat_batch(
    batch_size=8, num_epochs=None, num_steps=None)
for i, b in zip(range(6), infinite_bs):
  print(i)
Explanation: The number of batches produced from the iteration can be controlled by the (num_epochs, num_steps,
drop_remainder) combination:
If both num_epochs and num_steps are None, the shuffle-repeat process continues forever.
End of explanation
print('# batches w/ drop_remainder=False')
print(len(list(client_dataset.shuffle_repeat_batch(batch_size=8, num_epochs=1, num_steps=None))))
print('# batches w/ drop_remainder=True')
print(len(list(client_dataset.shuffle_repeat_batch(batch_size=8, num_epochs=1, num_steps=None, drop_remainder=True))))
Explanation: If num_epochs is set and num_steps is None, as few batches as needed to go
over the dataset this many passes are produced. Further,
If drop_remainder is False (the default), the final batch is filled with
    additionally sampled examples to contain batch_size examples.
If drop_remainder is True, the final batch is dropped if it contains fewer
    than batch_size examples. This may result in examples being skipped when
    num_epochs=1.
End of explanation
print('# batches w/ num_steps set and drop_remainder=True')
print(len(list(client_dataset.shuffle_repeat_batch(batch_size=8, num_epochs=None, num_steps=3, drop_remainder=True))))
Explanation: If num_steps is set and num_steps is None, exactly this many batches are
produced. drop_remainder has no effect in this case.
End of explanation
print('# batches w/ num_epochs and num_steps set')
print(len(list(client_dataset.shuffle_repeat_batch(batch_size=8, num_epochs=1, num_steps=6))))
Explanation: If both num_epochs and num_steps are set, the fewer number of batches
between the two conditions are produced.
End of explanation
# Random shuffling.
print(list(client_dataset.shuffle_repeat_batch(batch_size=2, seed=None))[0])
# Fixed shuffling.
print(list(client_dataset.shuffle_repeat_batch(batch_size=2, seed=0))[0])
Explanation: If reproducible iteration order is desired, a fixed seed can be used. When
seed is None, repeated iteration over the same object may produce batches in a
different order.
End of explanation
# A preprocessing function should return a new dict of examples instead of
# modifying its input.
def parity_feature(examples):
  return {'z': examples['y'] % 2, **examples}
# preprocess_batch returns a new FederatedData object that has one more
# preprocessing step at the very end than the original.
train_fd_z = train_fd.preprocess_batch(parity_feature)
client_id = b'105f96df763d4ddf:ALL_S_WELL_THAT_ENDS_WELL_GUIDERIUS_AND_ARVIRAGUS'
next(iter(train_fd_z.get_client(client_id).padded_batch(batch_size=4)))
Explanation: Preprocessing
Going from the raw features to features in a batch of examples ready for use in training/evalution often requires a few steps of preprocessing. Sometimes, we also want to add new preprocessing to an existing FederatedData.
Before going into the details of preprocessing, please note once again that all dataset related operations should be implemented in standard NumPy, not in JAX.
Preprocessing a batch of examples
The easiest way to add an additional preprocessing step is by appending a function that transforms a batch of examples to a FederatedData's list of preprocessing transformations on batches.
Below, we add a new z feature that stores the parity of y for our Shakespeare examples.
End of explanation
# Load unpreprocessed data.
raw_fd = fedjax.datasets.shakespeare.load_split('train')
raw_fd.get_client(client_id).all_examples()
Explanation: Preprocessing at the client level
Sometimes we also need to do some preprocessing for the entire client dataset. For example, in the Shakespeare dataset, the raw features are just text strings, so we need to turn them into a NumPy array of chunks of integer labels just to be able to meaningfully batch them at all.
In most circumstances, adding a preprocessing step at the client level is unnecessary, and should be avoided, because the new preprocessing step added by preprocess_client is inserted into the middle of a chain of steps, just before all the batch level preprocessing steps registered by preprocess_batch. If not done carefully, a custom client level preprocessing can easily break the preprocessing chain.
Nevertheless, here's an example of client level processing for certain rare cases.
End of explanation
# We don't actually need client_id, but `FederatedData` supplies it so that
# different processing based on clients can be done.
def truncate_and_cast(client_id, examples, max_length=10):
  labels = list(b''.join(examples['snippets']))
  num_sequences = (len(labels) + max_length - 1) // max_length
  padded = np.zeros((num_sequences, max_length), dtype=np.int32)
  for i in range(num_sequences):
    chars = labels[i * max_length:(i + 1) * max_length]
    padded[i, :len(chars)] = chars
  return {'snippets': padded}
partial_fd = raw_fd.preprocess_client(truncate_and_cast)
partial_fd.get_client(client_id).all_examples()
Explanation: The actual client level preprocessing in the Shakespeare dataset is a bit involved, so let's do something simpler: we shall join all the snippets, and then split and pad the integer byte values into bounded length sequences.
End of explanation
def snippets_to_xy(examples):
  snippets = examples['snippets']
  return {'x': snippets[:, :-1], 'y': snippets[:, 1:]}
partial_fd.preprocess_batch(snippets_to_xy).get_client(client_id).all_examples()
Explanation: Now, we can add another batch level preprocessor to produce x and y labels.
End of explanation
# Obtain MNIST dataset from tensorflow and convert to numpy format.
import tensorflow_datasets as tfds
(ds_train, ds_test) = tfds.load('mnist',
                                split=['train', 'test'],
                                shuffle_files=True,
                                as_supervised=True,
                                with_info=False)
features, labels = list(ds_train.batch(60000).as_numpy_iterator())[0]
print('features shape', features.shape)
print('labels shape', labels.shape)
# Randomly split dataset into 100 clients and load them to a dictionary.
indices = np.random.randint(100, size=60000)
client_id_to_dataset_mapping = {}
for i in range(100):
  client_id_to_dataset_mapping[i] = {'x': features[indices==i, :, : , :],
                                     'y': labels[indices==i]}
# Create fedjax.InMemoryDataset.
iid_mnist_federated_data = fedjax.InMemoryFederatedData(
    client_id_to_dataset_mapping)
print('number of clients in iid_mnist_data',
      iid_mnist_federated_data.num_clients())
Explanation: In memory federated datasets
In many scenarios, it is desirable to create a small custom federated dataset from a collection of NumPy arrays for quick experimentation. FedJAX provides fedjax.InMemoryFederatedData to create small custom datasets. fedjax.InMemoryFederatedData takes a dictionary of numpy examples keyed by client id and creates a fedjax.FederatedData that is compatible with the rest of the library. We illustrate this below with a simple example.
End of explanation
<END_TASK>