ValueError: The `input_spec` is a collections.abc.Mapping (e.g., a dict), so it must contain an entry with key `'x'`, representing the input(s)

Question

I'm testing this tutorial with non-IID distribution for federated learning: https://www.tensorflow.org/federated/tutorials/tff_for_federated_learning_research_compression, and using tff.simulation.datasets.build_single_label_dataset() as a way to produce a non-IID distribution for the dataset.

But I faced an error regarding keras values in the input_spec.

The code:

emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data(
    only_digits=False)

emnist_train1 = tff.simulation.datasets.build_single_label_dataset(
  emnist_train.create_tf_dataset_from_all_clients(),
  label_key='label', desired_label=1)

MAX_CLIENT_DATASET_SIZE = 418

CLIENT_EPOCHS_PER_ROUND = 1
CLIENT_BATCH_SIZE = 20
TEST_BATCH_SIZE = 500

def reshape_emnist_element(element):
  return (tf.expand_dims(element['pixels'], axis=-1), element['label'])

def preprocess_train_dataset(dataset):
  return (dataset
          .shuffle(buffer_size=MAX_CLIENT_DATASET_SIZE)
          .repeat(CLIENT_EPOCHS_PER_ROUND)
          .batch(CLIENT_BATCH_SIZE, drop_remainder=False)
          .map(reshape_emnist_element))

emnist_train1 = preprocess_train_dataset(emnist_train1)

import random
NUM_CLIENTS = 100

client_datasets = [
   emnist_train1.take(random.randint(1, CLIENT_BATCH_SIZE))
   for _ in range(NUM_CLIENTS)
]

emnist_train1 = client_datasets

def create_original_fedavg_cnn_model(only_digits=True):
  """The CNN model used in https://arxiv.org/abs/1602.05629."""
  data_format = 'channels_last'

  max_pool = functools.partial(
      tf.keras.layers.MaxPooling2D,
      pool_size=(2, 2),
      padding='same',
      data_format=data_format)
  conv2d = functools.partial(
      tf.keras.layers.Conv2D,
      kernel_size=5,
      padding='same',
      data_format=data_format,
      activation=tf.nn.relu)

  model = tf.keras.models.Sequential([
      tf.keras.layers.InputLayer(input_shape=(28, 28, 1)),
      conv2d(filters=32),
      max_pool(),
      conv2d(filters=64),
      max_pool(),
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(512, activation=tf.nn.relu),
      tf.keras.layers.Dense(10 if only_digits else 62),
      tf.keras.layers.Softmax(),
  ])

  return model

input_spec = emnist_train.create_tf_dataset_for_client(
    emnist_train.client_ids[0]).element_spec

def tff_model_fn():
  keras_model = create_original_fedavg_cnn_model()
  return tff.learning.from_keras_model(
      keras_model=keras_model,
      input_spec=input_spec,
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])

federated_averaging = tff.learning.build_federated_averaging_process(
    model_fn=tff_model_fn,
    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02),
    server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0))
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-11-7661a3b7dc92> in <module>()
      3     model_fn=tff_model_fn,
      4     client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02),
----> 5     server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0))

6 frames
/usr/local/lib/python3.7/dist-packages/tensorflow_federated/python/learning/keras_utils.py in from_keras_model(keras_model, loss, input_spec, loss_weights, metrics)
    176           'The `input_spec` is a collections.abc.Mapping (e.g., a dict), so it '
    177           'must contain an entry with key `\'{}\'`, representing the input(s) '
--> 178           'to the Keras model.'.format(model_lib.MODEL_ARG_NAME))
    179     if model_lib.MODEL_LABEL_NAME not in input_spec:
    180       raise ValueError(

ValueError: The input_spec is a collections.abc.Mapping (e.g., a dict), so it must contain an entry with key 'x', representing the input(s) to the Keras model.

What does that mean? How can I solve it?