# Copyright 2018, The TensorFlow Federated Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Example showing how to run a multi-machine simulation.

In order to run this example, you must have a running instance of the
Executor Service, either locally or on Kubernetes.

The model trains EMNIST for a small number of rounds, but uses a RemoteExecutor
to distribute the work to the ExecutorService.
"""

import collections
import warnings

from absl import app
from absl import flags
import grpc
import numpy as np
import tensorflow as tf
import tensorflow_federated as tff

FLAGS = flags.FLAGS

flags.DEFINE_string('host', None, 'The host to connect to.')
flags.mark_flag_as_required('host')
flags.DEFINE_string('port', '8000', 'The port to connect to.')
flags.DEFINE_integer('n_clients', 10, 'Number of clients.')
flags.DEFINE_integer('n_rounds', 3, 'Number of rounds.')


def preprocess(dataset):

  def element_fn(element):
    return collections.OrderedDict([
        ('x', tf.reshape(element['pixels'], [-1])),
        ('y', tf.reshape(element['label'], [1])),
    ])

  return dataset.repeat(NUM_EPOCHS).map(element_fn).batch(BATCH_SIZE)


def make_federated_data(client_data, client_ids):
  return [
      preprocess(client_data.create_tf_dataset_for_client(x))
      for x in client_ids
  ]


NUM_EPOCHS = 10
BATCH_SIZE = 20


def make_remote_executor(inferred_cardinalities):
  """Make remote executor."""

  def create_worker_stack(ex):
    ex = tff.framework.ThreadDelegatingExecutor(ex)
    return tff.framework.ReferenceResolvingExecutor(ex)

  client_ex = []
  num_clients = inferred_cardinalities.get(tff.CLIENTS, None)
  if num_clients:
    print('Inferred that there are {} clients'.format(num_clients))
  else:
    print('No CLIENTS placement provided')

  for _ in range(num_clients or 0):
    channel = grpc.insecure_channel('{}:{}'.format(FLAGS.host, FLAGS.port))
    remote_ex = tff.framework.RemoteExecutor(channel)
    worker_stack = create_worker_stack(remote_ex)
    client_ex.append(worker_stack)

  federating_strategy_factory = tff.framework.FederatedResolvingStrategy.factory(
      {
          tff.SERVER: create_worker_stack(tff.framework.EagerTFExecutor()),
          tff.CLIENTS: client_ex,
      })
  unplaced_ex = create_worker_stack(tff.framework.EagerTFExecutor())
  federating_ex = tff.framework.FederatingExecutor(federating_strategy_factory,
                                                   unplaced_ex)
  return tff.framework.ReferenceResolvingExecutor(federating_ex)


def main(argv):
  if len(argv) > 1:
    raise app.UsageError('Too many command-line arguments.')

  warnings.simplefilter('ignore')

  np.random.seed(0)

  emnist_train, _ = tff.simulation.datasets.emnist.load_data()

  sample_clients = emnist_train.client_ids[0:FLAGS.n_clients]

  federated_train_data = make_federated_data(emnist_train, sample_clients)

  example_dataset = emnist_train.create_tf_dataset_for_client(
      emnist_train.client_ids[0])

  preprocessed_example_dataset = preprocess(example_dataset)
  input_spec = preprocessed_example_dataset.element_spec

  def model_fn():
    model = tf.keras.models.Sequential([
        tf.keras.layers.Input(shape=(784,)),
        tf.keras.layers.Dense(10, kernel_initializer='zeros'),
        tf.keras.layers.Softmax(),
    ])
    return tff.learning.from_keras_model(
        model,
        input_spec=input_spec,
        loss=tf.keras.losses.SparseCategoricalCrossentropy(),
        metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])

  iterative_process = tff.learning.build_federated_averaging_process(
      model_fn,
      client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02))

  factory = tff.framework.ResourceManagingExecutorFactory(make_remote_executor)
  context = tff.framework.ExecutionContext(factory)
  tff.framework.set_default_context(context)

  state = iterative_process.initialize()

  state, metrics = iterative_process.next(state, federated_train_data)
  print('round  1, metrics={}'.format(metrics))

  for round_num in range(2, FLAGS.n_rounds + 1):
    state, metrics = iterative_process.next(state, federated_train_data)
    print('round {:2d}, metrics={}'.format(round_num, metrics))


if __name__ == '__main__':
  app.run(main)

tensorflow_federated/python/examples/remote_execution/remote_executor_example.py
# Copyright 2018, The TensorFlow Federated Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Example showing how to run a multi-machine simulation.

In order to run this example, you must have a running instance of the
Executor Service, either locally or on Kubernetes.

The model trains EMNIST for a small number of rounds, but uses a RemoteExecutor
to distribute the work to the ExecutorService.
"""

import collections
import warnings

from absl import app
from absl import flags
import grpc
import numpy as np
import tensorflow as tf
import tensorflow_federated as tff

FLAGS = flags.FLAGS

flags.DEFINE_string('host', None, 'The host to connect to.')
flags.mark_flag_as_required('host')
flags.DEFINE_string('port', '8000', 'The port to connect to.')
flags.DEFINE_integer('n_clients', 10, 'Number of clients.')
flags.DEFINE_integer('n_rounds', 3, 'Number of rounds.')


def preprocess(dataset):

  def element_fn(element):
    return collections.OrderedDict([
        ('x', tf.reshape(element['pixels'], [-1])),
        ('y', tf.reshape(element['label'], [1])),
    ])

  return dataset.repeat(NUM_EPOCHS).map(element_fn).batch(BATCH_SIZE)


def make_federated_data(client_data, client_ids):
  return [
      preprocess(client_data.create_tf_dataset_for_client(x))
      for x in client_ids
  ]


NUM_EPOCHS = 10
BATCH_SIZE = 20


def make_remote_executor(inferred_cardinalities):
  """Make remote executor."""

  def create_worker_stack(ex):
    ex = tff.framework.ThreadDelegatingExecutor(ex)
    return tff.framework.ReferenceResolvingExecutor(ex)

  client_ex = []
  num_clients = inferred_cardinalities.get(tff.CLIENTS, None)
  if num_clients:
    print('Inferred that there are {} clients'.format(num_clients))
  else:
    print('No CLIENTS placement provided')

  for _ in range(num_clients or 0):
    channel = grpc.insecure_channel('{}:{}'.format(FLAGS.host, FLAGS.port))
    remote_ex = tff.framework.RemoteExecutor(channel)
    worker_stack = create_worker_stack(remote_ex)
    client_ex.append(worker_stack)

  federating_strategy_factory = tff.framework.FederatedResolvingStrategy.factory(
      {
          tff.SERVER: create_worker_stack(tff.framework.EagerTFExecutor()),
          tff.CLIENTS: client_ex,
      })
  unplaced_ex = create_worker_stack(tff.framework.EagerTFExecutor())
  federating_ex = tff.framework.FederatingExecutor(federating_strategy_factory,
                                                   unplaced_ex)
  return tff.framework.ReferenceResolvingExecutor(federating_ex)


def main(argv):
  if len(argv) > 1:
    raise app.UsageError('Too many command-line arguments.')

  warnings.simplefilter('ignore')

  np.random.seed(0)

  emnist_train, _ = tff.simulation.datasets.emnist.load_data()

  sample_clients = emnist_train.client_ids[0:FLAGS.n_clients]

  federated_train_data = make_federated_data(emnist_train, sample_clients)

  example_dataset = emnist_train.create_tf_dataset_for_client(
      emnist_train.client_ids[0])

  preprocessed_example_dataset = preprocess(example_dataset)
  input_spec = preprocessed_example_dataset.element_spec

  def model_fn():
    model = tf.keras.models.Sequential([
        tf.keras.layers.InputLayer(input_shape=(784,)),
        tf.keras.layers.Dense(10, kernel_initializer='zeros'),
        tf.keras.layers.Softmax(),
    ])
    return tff.learning.from_keras_model(
        model,
        input_spec=input_spec,
        loss=tf.keras.losses.SparseCategoricalCrossentropy(),
        metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])

  iterative_process = tff.learning.build_federated_averaging_process(
      model_fn,
      client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02))

  factory = tff.framework.ResourceManagingExecutorFactory(make_remote_executor)
  context = tff.framework.ExecutionContext(factory)
  tff.framework.set_default_context(context)

  state = iterative_process.initialize()

  state, metrics = iterative_process.next(state, federated_train_data)
  print('round  1, metrics={}'.format(metrics))

  for round_num in range(2, FLAGS.n_rounds + 1):
    state, metrics = iterative_process.next(state, federated_train_data)
    print('round {:2d}, metrics={}'.format(round_num, metrics))


if __name__ == '__main__':
  app.run(main)

tensorflow_federated/python/examples/remote_execution/remote_executor_example.py
diff --git a/docs/openmined2020/openmined_conference_2020.ipynb b/docs/openmined2020/openmined_conference_2020.ipynb
index 16ce40a125..95e00dfb10 100644
--- a/docs/openmined2020/openmined_conference_2020.ipynb
+++ b/docs/openmined2020/openmined_conference_2020.ipynb
@@ -499,7 +499,7 @@
       "source": [
         "def create_keras_model():\n",
         "  return tf.keras.models.Sequential([\n",
-        "      tf.keras.layers.Input(shape=(784,)),\n",
+        "      tf.keras.layers.InputLayer(input_shape=(784,)),\n",
         "      tf.keras.layers.Dense(10, kernel_initializer='zeros'),\n",
         "      tf.keras.layers.Softmax(),\n",
         "  ])"
@@ -1083,7 +1083,7 @@
       "source": [
         "def create_keras_model():\n",
         "  return tf.keras.models.Sequential([\n",
-        "      tf.keras.layers.Input(shape=(784,)),\n",
+        "      tf.keras.layers.InputLayer(input_shape=(784,)),\n",
         "      tf.keras.layers.Dense(10, kernel_initializer='zeros'),\n",
         "      tf.keras.layers.Softmax(),\n",
         "  ])"
diff --git a/docs/tutorials/federated_learning_for_image_classification.ipynb b/docs/tutorials/federated_learning_for_image_classification.ipynb
index eed5fc4404..de7f86a3a8 100644
--- a/docs/tutorials/federated_learning_for_image_classification.ipynb
+++ b/docs/tutorials/federated_learning_for_image_classification.ipynb
@@ -943,7 +943,7 @@
       "source": [
         "def create_keras_model():\n",
         "  return tf.keras.models.Sequential([\n",
-        "      tf.keras.layers.Input(shape=(784,)),\n",
+        "      tf.keras.layers.InputLayer(input_shape=(784,)),\n",
         "      tf.keras.layers.Dense(10, kernel_initializer='zeros'),\n",
         "      tf.keras.layers.Softmax(),\n",
         "  ])"
diff --git a/docs/tutorials/high_performance_simulation_with_kubernetes.ipynb b/docs/tutorials/high_performance_simulation_with_kubernetes.ipynb
index db9495a670..33e2f98c61 100644
--- a/docs/tutorials/high_performance_simulation_with_kubernetes.ipynb
+++ b/docs/tutorials/high_performance_simulation_with_kubernetes.ipynb
@@ -197,7 +197,7 @@
         "\n",
         "def model_fn():\n",
         "  model = tf.keras.models.Sequential([\n",
-        "      tf.keras.layers.Input(shape=(784,)),\n",
+        "      tf.keras.layers.InputLayer(input_shape=(784,)),\n",
         "      tf.keras.layers.Dense(units=10, kernel_initializer='zeros'),\n",
         "      tf.keras.layers.Softmax(),\n",
         "  ])\n",
diff --git a/docs/tutorials/simulations.ipynb b/docs/tutorials/simulations.ipynb
index 4cb2651621..f221410436 100644
--- a/docs/tutorials/simulations.ipynb
+++ b/docs/tutorials/simulations.ipynb
@@ -153,7 +153,7 @@
         "\n",
         "def model_fn():\n",
         "  model = tf.keras.models.Sequential([\n",
-        "      tf.keras.layers.Input(shape=(784,)),\n",
+        "      tf.keras.layers.InputLayer(input_shape=(784,)),\n",
         "      tf.keras.layers.Dense(units=10, kernel_initializer='zeros'),\n",
         "      tf.keras.layers.Softmax(),\n",
         "  ])\n",
diff --git a/docs/tutorials/working_with_client_data.ipynb b/docs/tutorials/working_with_client_data.ipynb
index 05dd20004b..8db413d150 100644
--- a/docs/tutorials/working_with_client_data.ipynb
+++ b/docs/tutorials/working_with_client_data.ipynb
@@ -271,7 +271,7 @@
       "source": [
         "def model_fn():\n",
         "  model = tf.keras.models.Sequential([\n",
-        "      tf.keras.layers.Input(shape=(784,)),\n",
+        "      tf.keras.layers.InputLayer(input_shape=(784,)),\n",
         "      tf.keras.layers.Dense(10, kernel_initializer='zeros'),\n",
         "  ])\n",
         "  return tff.learning.from_keras_model(\n",
diff --git a/tensorflow_federated/python/examples/remote_execution/remote_executor_example.py b/tensorflow_federated/python/examples/remote_execution/remote_executor_example.py
index f7962471d8..a54d1f5cef 100644
--- a/tensorflow_federated/python/examples/remote_execution/remote_executor_example.py
+++ b/tensorflow_federated/python/examples/remote_execution/remote_executor_example.py
@@ -114,7 +114,7 @@ def main(argv):
 
   def model_fn():
     model = tf.keras.models.Sequential([
-        tf.keras.layers.Input(shape=(784,)),
+        tf.keras.layers.InputLayer(input_shape=(784,)),
         tf.keras.layers.Dense(10, kernel_initializer='zeros'),
         tf.keras.layers.Softmax(),
     ])
diff --git a/tensorflow_federated/python/learning/federated_evaluation_test.py b/tensorflow_federated/python/learning/federated_evaluation_test.py
index 855b44a211..a30d472712 100644
--- a/tensorflow_federated/python/learning/federated_evaluation_test.py
+++ b/tensorflow_federated/python/learning/federated_evaluation_test.py
@@ -144,7 +144,7 @@ def aggregate_metrics(client_metrics):
 
 def _model_fn_from_keras():
   keras_model = tf.keras.Sequential([
-      tf.keras.layers.Input(shape=(1,)),
+      tf.keras.layers.InputLayer(input_shape=(1,)),
       tf.keras.layers.Dense(
           1,
           kernel_initializer='ones',
diff --git a/tensorflow_federated/python/learning/keras_utils_test.py b/tensorflow_federated/python/learning/keras_utils_test.py
index 42e5c2fb7d..e14b762680 100644
--- a/tensorflow_federated/python/learning/keras_utils_test.py
+++ b/tensorflow_federated/python/learning/keras_utils_test.py
@@ -379,7 +379,7 @@ def call(self, y_true, y_pred):
 
   def test_tff_model_type_spec_from_keras_model_unspecified_sequence_len(self):
     keras_model = tf.keras.Sequential([
-        tf.keras.layers.Input(shape=(None,)),
+        tf.keras.layers.InputLayer(input_shape=(None,)),
         tf.keras.layers.Embedding(input_dim=10, output_dim=10),
         tf.keras.layers.LSTM(1)
     ])
diff --git a/tensorflow_federated/python/tests/backend_test.py b/tensorflow_federated/python/tests/backend_test.py
index 480c8c6c46..056926fc8d 100644
--- a/tensorflow_federated/python/tests/backend_test.py
+++ b/tensorflow_federated/python/tests/backend_test.py
@@ -390,7 +390,7 @@ def client_data(n):
 
     def model_fn():
       model = tf.keras.models.Sequential([
-          tf.keras.layers.Input(shape=(input_dim,)),
+          tf.keras.layers.InputLayer(input_shape=(input_dim,)),
           tf.keras.layers.Dense(
               50, activity_regularizer=tf.keras.regularizers.l1(.1)),
           tf.keras.layers.Dense(input_dim, activation='sigmoid'),

Keras model in federated_learning_for_image_classification.ipynb throws warning
**Describe the bug**
Keras Sequential Model in [federated_learning_for_image_classification.ipynb](https://github.com/tensorflow/federated/blob/master/docs/tutorials/federated_learning_for_image_classification.ipynb) throws warning.
The model in the notebook is
```python
def create_keras_model():
  return tf.keras.models.Sequential([
      tf.keras.layers.Input(shape=(784,)),
      tf.keras.layers.Dense(10, kernel_initializer='zeros'),
      tf.keras.layers.Softmax(),
  ])
```
Warning thrown:
```python
WARNING:tensorflow:Please add `keras.layers.InputLayer` instead of `keras.Input` to Sequential model. `keras.Input` is intended to be used by Functional model.
<tensorflow.python.keras.engine.sequential.Sequential at 0x7f66178a46d0>
```

Easily fixed using the correct layer type:
```python
def create_keras_model():
  return tf.keras.models.Sequential([
      tf.keras.layers.InputLayer(input_shape=(784,)),
      tf.keras.layers.Dense(10, kernel_initializer='zeros'),
      tf.keras.layers.Softmax(),
  ])
```

[colab](https://colab.research.google.com/drive/1LFgBiu9xUa-k92IW24fiSX_kVp7lb0SB?usp=sharing) notebook that reproduces the bug.

