We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Modify elasticdl.layers.Embedding arguments and constructor according to design doc
According to [design doc](https://github.com/wangkuiyi/elasticdl/blob/develop/elasticdl/doc/distributed_embedding_layer_design.md#elasticdllayersembedding):

```
__init__(
    output_dim,
    embeddings_initializer='uniform',
    mask_zero=False,
    input_length=None,
    combiner=None,
)
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `elasticdl/python/elasticdl/layers/embedding.py`
Content:
```
1 import tensorflow as tf
2 from tensorflow.python.keras.utils import tf_utils
3 
4 
5 class Embedding(tf.keras.layers.Layer):
6     """
7     Input: indexes for the embedding entries
8            shape is (batch_size, input_length)
9     Output: Corresponding embedding vectors of the input indexes
10             shape is (batch_size, input_length, embedding_dim)
11     Arguments:
12       embedding_dim: the dimension of the embedding vector
13       embedding_initializer: Initializer for embedding table
14     """
15 
16     def __init__(self, embedding_dim, embedding_initializer="uniform"):
17         super(Embedding, self).__init__()
18         self.embedding_dim = embedding_dim
19         self.embedding_initializer = embedding_initializer
20         self.tape = None
21         self.worker = None
22         self.bet_ids_pair = []
23 
24     @tf_utils.shape_type_conversion
25     def compute_output_shape(self, input_shape):
26         return input_shape + (self.embedding_dim,)
27 
28     @property
29     def name(self):
30         return self._name
31 
32     @staticmethod
33     def get_key(name_list):
34         return "-".join(map(str, name_list))
35 
36     def lookup_embedding(self, unique_ids):
37         batch_embedding = self.worker.embedding_lookup(
38             unique_ids, self._name, self.embedding_initializer
39         )
40         return batch_embedding
41 
42     def call(self, input):
43         ids = tf.convert_to_tensor(input, name="embedding_ids")
44         flat_ids = tf.reshape(ids, [-1])
45         unique_ids, idx = tf.unique(flat_ids)
46         batch_embedding_tensor = tf.py_function(
47             self.lookup_embedding, inp=[unique_ids], Tout=tf.float32
48         )
49         if self.tape:
50             # tape.watch works with eager mode only
51             if not tf.executing_eagerly():
52                 raise RuntimeError("tape.watch only works with eager mode")
53             self.tape.watch(batch_embedding_tensor)
54             self.bet_ids_pair.append((batch_embedding_tensor, unique_ids))
55         outputs = tf.gather(batch_embedding_tensor, idx)
56         outputs = tf.reshape(
57             outputs, ids.get_shape().concatenate(self.embedding_dim)
58         )
59         return outputs
60 
61     def reset(self):
62         self.bet_ids_pair = []
63         self.tape = None
64 
65     def set_tape(self, tape):
66         self.tape = tape
67 
68     def set_worker(self, worker):
69         self.worker = worker
70
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/elasticdl/python/elasticdl/layers/embedding.py b/elasticdl/python/elasticdl/layers/embedding.py
--- a/elasticdl/python/elasticdl/layers/embedding.py
+++ b/elasticdl/python/elasticdl/layers/embedding.py
@@ -6,24 +6,77 @@
     """
     Input: indexes for the embedding entries
            shape is (batch_size, input_length)
-    Output: Corresponding embedding vectors of the input indexes
-            shape is (batch_size, input_length, embedding_dim)
+    Output:
+      corresponding (combined) embeddings with a shape of
+      (batch_size, input_length, output_dim) if combiner is None
+      (batch_size, output_dim) if combiner is not None
     Arguments:
-      embedding_dim: the dimension of the embedding vector
+      output_dim: the dimension of the embedding vector
       embedding_initializer: Initializer for embedding table
+      mask_zero: Whether or not the input value 0 is a special "padding"
+        value that should be masked out.
+      input_length: Length of input sequences, when it is constant.
+        This argument is required if you are going to connect
+        `Flatten` then `Dense` layers upstream
+        (without it, the shape of the dense outputs cannot be computed).
+      combiner: A string specifying the reduction op or None if not used.
+        "mean", "sqrtn" and "sum" are supported for the reduction op.
+    TODO: support mask_zero
+    TODO: support combiner
+    TODO: support sparse input
     """
 
-    def __init__(self, embedding_dim, embedding_initializer="uniform"):
-        super(Embedding, self).__init__()
-        self.embedding_dim = embedding_dim
+    def __init__(
+        self,
+        output_dim,
+        embedding_initializer="uniform",
+        mask_zero=False,
+        input_length=None,
+        combiner=None,
+        **kwargs
+    ):
+        if "input_shape" not in kwargs and input_length:
+            kwargs["input_shape"] = (input_length,)
+        super(Embedding, self).__init__(**kwargs)
+
+        self.output_dim = output_dim
         self.embedding_initializer = embedding_initializer
+        self.mask_zero = mask_zero
+        self.input_length = input_length
+        self.combiner = combiner
         self.tape = None
         self.worker = None
         self.bet_ids_pair = []
 
     @tf_utils.shape_type_conversion
     def compute_output_shape(self, input_shape):
-        return input_shape + (self.embedding_dim,)
+        # this function is taken from
+        # tf.keras.layers.Embedding.compute_output_shape
+        # https://github.com/tensorflow/tensorflow/blob/3f3c728bf80e0fd6653744318cbbfe1454c6ddca/tensorflow/python/keras/layers/embeddings.py#L156
+        if self.input_length is None:
+            return input_shape + (self.output_dim,)
+        else:
+            if isinstance(self.input_length, (list, tuple)):
+                in_lens = list(self.input_length)
+            else:
+                in_lens = [self.input_length]
+            if len(in_lens) != len(input_shape) - 1:
+                raise ValueError(
+                    '"input_length" is %s, '
+                    "but received input has shape %s"
+                    % (str(self.input_length), str(input_shape))
+                )
+            else:
+                for i, (s1, s2) in enumerate(zip(in_lens, input_shape[1:])):
+                    if s1 is not None and s2 is not None and s1 != s2:
+                        raise ValueError(
+                            '"input_length" is %s, '
+                            "but received input has shape %s"
+                            % (str(self.input_length), str(input_shape))
+                        )
+                    elif s1 is None:
+                        in_lens[i] = s2
+            return (input_shape[0],) + tuple(in_lens) + (self.output_dim,)
 
     @property
     def name(self):
@@ -54,7 +107,7 @@
             self.bet_ids_pair.append((batch_embedding_tensor, unique_ids))
         outputs = tf.gather(batch_embedding_tensor, idx)
         outputs = tf.reshape(
-            outputs, ids.get_shape().concatenate(self.embedding_dim)
+            outputs, ids.get_shape().concatenate(self.output_dim)
         )
         return outputs

{"golden_diff": "diff --git a/elasticdl/python/elasticdl/layers/embedding.py b/elasticdl/python/elasticdl/layers/embedding.py\n--- a/elasticdl/python/elasticdl/layers/embedding.py\n+++ b/elasticdl/python/elasticdl/layers/embedding.py\n@@ -6,24 +6,77 @@\n     \"\"\"\n     Input: indexes for the embedding entries\n            shape is (batch_size, input_length)\n-    Output: Corresponding embedding vectors of the input indexes\n-            shape is (batch_size, input_length, embedding_dim)\n+    Output:\n+      corresponding (combined) embeddings with a shape of\n+      (batch_size, input_length, output_dim) if combiner is None\n+      (batch_size, output_dim) if combiner is not None\n     Arguments:\n-      embedding_dim: the dimension of the embedding vector\n+      output_dim: the dimension of the embedding vector\n       embedding_initializer: Initializer for embedding table\n+      mask_zero: Whether or not the input value 0 is a special \"padding\"\n+        value that should be masked out.\n+      input_length: Length of input sequences, when it is constant.\n+        This argument is required if you are going to connect\n+        `Flatten` then `Dense` layers upstream\n+        (without it, the shape of the dense outputs cannot be computed).\n+      combiner: A string specifying the reduction op or None if not used.\n+        \"mean\", \"sqrtn\" and \"sum\" are supported for the reduction op.\n+    TODO: support mask_zero\n+    TODO: support combiner\n+    TODO: support sparse input\n     \"\"\"\n \n-    def __init__(self, embedding_dim, embedding_initializer=\"uniform\"):\n-        super(Embedding, self).__init__()\n-        self.embedding_dim = embedding_dim\n+    def __init__(\n+        self,\n+        output_dim,\n+        embedding_initializer=\"uniform\",\n+        mask_zero=False,\n+        input_length=None,\n+        combiner=None,\n+        **kwargs\n+    ):\n+        if \"input_shape\" not in kwargs and input_length:\n+            kwargs[\"input_shape\"] = (input_length,)\n+        super(Embedding, self).__init__(**kwargs)\n+\n+        self.output_dim = output_dim\n         self.embedding_initializer = embedding_initializer\n+        self.mask_zero = mask_zero\n+        self.input_length = input_length\n+        self.combiner = combiner\n         self.tape = None\n         self.worker = None\n         self.bet_ids_pair = []\n \n     @tf_utils.shape_type_conversion\n     def compute_output_shape(self, input_shape):\n-        return input_shape + (self.embedding_dim,)\n+        # this function is taken from\n+        # tf.keras.layers.Embedding.compute_output_shape\n+        # https://github.com/tensorflow/tensorflow/blob/3f3c728bf80e0fd6653744318cbbfe1454c6ddca/tensorflow/python/keras/layers/embeddings.py#L156\n+        if self.input_length is None:\n+            return input_shape + (self.output_dim,)\n+        else:\n+            if isinstance(self.input_length, (list, tuple)):\n+                in_lens = list(self.input_length)\n+            else:\n+                in_lens = [self.input_length]\n+            if len(in_lens) != len(input_shape) - 1:\n+                raise ValueError(\n+                    '\"input_length\" is %s, '\n+                    \"but received input has shape %s\"\n+                    % (str(self.input_length), str(input_shape))\n+                )\n+            else:\n+                for i, (s1, s2) in enumerate(zip(in_lens, input_shape[1:])):\n+                    if s1 is not None and s2 is not None and s1 != s2:\n+                        raise ValueError(\n+                            '\"input_length\" is %s, '\n+                            \"but received input has shape %s\"\n+                            % (str(self.input_length), str(input_shape))\n+                        )\n+                    elif s1 is None:\n+                        in_lens[i] = s2\n+            return (input_shape[0],) + tuple(in_lens) + (self.output_dim,)\n \n     @property\n     def name(self):\n@@ -54,7 +107,7 @@\n             self.bet_ids_pair.append((batch_embedding_tensor, unique_ids))\n         outputs = tf.gather(batch_embedding_tensor, idx)\n         outputs = tf.reshape(\n-            outputs, ids.get_shape().concatenate(self.embedding_dim)\n+            outputs, ids.get_shape().concatenate(self.output_dim)\n         )\n         return outputs\n", "issue": "Modify elasticdl.layers.Embedding arguments and constructor according to design doc\nAccording to [design doc](https://github.com/wangkuiyi/elasticdl/blob/develop/elasticdl/doc/distributed_embedding_layer_design.md#elasticdllayersembedding):\r\n\r\n```\r\n__init__(\r\n    output_dim,\r\n    embeddings_initializer='uniform',\r\n    mask_zero=False,\r\n    input_length=None,\r\n    combiner=None,\r\n)\r\n```\n", "before_files": [{"content": "import tensorflow as tf\nfrom tensorflow.python.keras.utils import tf_utils\n\n\nclass Embedding(tf.keras.layers.Layer):\n    \"\"\"\n    Input: indexes for the embedding entries\n           shape is (batch_size, input_length)\n    Output: Corresponding embedding vectors of the input indexes\n            shape is (batch_size, input_length, embedding_dim)\n    Arguments:\n      embedding_dim: the dimension of the embedding vector\n      embedding_initializer: Initializer for embedding table\n    \"\"\"\n\n    def __init__(self, embedding_dim, embedding_initializer=\"uniform\"):\n        super(Embedding, self).__init__()\n        self.embedding_dim = embedding_dim\n        self.embedding_initializer = embedding_initializer\n        self.tape = None\n        self.worker = None\n        self.bet_ids_pair = []\n\n    @tf_utils.shape_type_conversion\n    def compute_output_shape(self, input_shape):\n        return input_shape + (self.embedding_dim,)\n\n    @property\n    def name(self):\n        return self._name\n\n    @staticmethod\n    def get_key(name_list):\n        return \"-\".join(map(str, name_list))\n\n    def lookup_embedding(self, unique_ids):\n        batch_embedding = self.worker.embedding_lookup(\n            unique_ids, self._name, self.embedding_initializer\n        )\n        return batch_embedding\n\n    def call(self, input):\n        ids = tf.convert_to_tensor(input, name=\"embedding_ids\")\n        flat_ids = tf.reshape(ids, [-1])\n        unique_ids, idx = tf.unique(flat_ids)\n        batch_embedding_tensor = tf.py_function(\n            self.lookup_embedding, inp=[unique_ids], Tout=tf.float32\n        )\n        if self.tape:\n            # tape.watch works with eager mode only\n            if not tf.executing_eagerly():\n                raise RuntimeError(\"tape.watch only works with eager mode\")\n            self.tape.watch(batch_embedding_tensor)\n            self.bet_ids_pair.append((batch_embedding_tensor, unique_ids))\n        outputs = tf.gather(batch_embedding_tensor, idx)\n        outputs = tf.reshape(\n            outputs, ids.get_shape().concatenate(self.embedding_dim)\n        )\n        return outputs\n\n    def reset(self):\n        self.bet_ids_pair = []\n        self.tape = None\n\n    def set_tape(self, tape):\n        self.tape = tape\n\n    def set_worker(self, worker):\n        self.worker = worker\n", "path": "elasticdl/python/elasticdl/layers/embedding.py"}], "after_files": [{"content": "import tensorflow as tf\nfrom tensorflow.python.keras.utils import tf_utils\n\n\nclass Embedding(tf.keras.layers.Layer):\n    \"\"\"\n    Input: indexes for the embedding entries\n           shape is (batch_size, input_length)\n    Output:\n      corresponding (combined) embeddings with a shape of\n      (batch_size, input_length, output_dim) if combiner is None\n      (batch_size, output_dim) if combiner is not None\n    Arguments:\n      output_dim: the dimension of the embedding vector\n      embedding_initializer: Initializer for embedding table\n      mask_zero: Whether or not the input value 0 is a special \"padding\"\n        value that should be masked out.\n      input_length: Length of input sequences, when it is constant.\n        This argument is required if you are going to connect\n        `Flatten` then `Dense` layers upstream\n        (without it, the shape of the dense outputs cannot be computed).\n      combiner: A string specifying the reduction op or None if not used.\n        \"mean\", \"sqrtn\" and \"sum\" are supported for the reduction op.\n    TODO: support mask_zero\n    TODO: support combiner\n    TODO: support sparse input\n    \"\"\"\n\n    def __init__(\n        self,\n        output_dim,\n        embedding_initializer=\"uniform\",\n        mask_zero=False,\n        input_length=None,\n        combiner=None,\n        **kwargs\n    ):\n        if \"input_shape\" not in kwargs and input_length:\n            kwargs[\"input_shape\"] = (input_length,)\n        super(Embedding, self).__init__(**kwargs)\n\n        self.output_dim = output_dim\n        self.embedding_initializer = embedding_initializer\n        self.mask_zero = mask_zero\n        self.input_length = input_length\n        self.combiner = combiner\n        self.tape = None\n        self.worker = None\n        self.bet_ids_pair = []\n\n    @tf_utils.shape_type_conversion\n    def compute_output_shape(self, input_shape):\n        # this function is taken from\n        # tf.keras.layers.Embedding.compute_output_shape\n        # https://github.com/tensorflow/tensorflow/blob/3f3c728bf80e0fd6653744318cbbfe1454c6ddca/tensorflow/python/keras/layers/embeddings.py#L156\n        if self.input_length is None:\n            return input_shape + (self.output_dim,)\n        else:\n            if isinstance(self.input_length, (list, tuple)):\n                in_lens = list(self.input_length)\n            else:\n                in_lens = [self.input_length]\n            if len(in_lens) != len(input_shape) - 1:\n                raise ValueError(\n                    '\"input_length\" is %s, '\n                    \"but received input has shape %s\"\n                    % (str(self.input_length), str(input_shape))\n                )\n            else:\n                for i, (s1, s2) in enumerate(zip(in_lens, input_shape[1:])):\n                    if s1 is not None and s2 is not None and s1 != s2:\n                        raise ValueError(\n                            '\"input_length\" is %s, '\n                            \"but received input has shape %s\"\n                            % (str(self.input_length), str(input_shape))\n                        )\n                    elif s1 is None:\n                        in_lens[i] = s2\n            return (input_shape[0],) + tuple(in_lens) + (self.output_dim,)\n\n    @property\n    def name(self):\n        return self._name\n\n    @staticmethod\n    def get_key(name_list):\n        return \"-\".join(map(str, name_list))\n\n    def lookup_embedding(self, unique_ids):\n        batch_embedding = self.worker.embedding_lookup(\n            unique_ids, self._name, self.embedding_initializer\n        )\n        return batch_embedding\n\n    def call(self, input):\n        ids = tf.convert_to_tensor(input, name=\"embedding_ids\")\n        flat_ids = tf.reshape(ids, [-1])\n        unique_ids, idx = tf.unique(flat_ids)\n        batch_embedding_tensor = tf.py_function(\n            self.lookup_embedding, inp=[unique_ids], Tout=tf.float32\n        )\n        if self.tape:\n            # tape.watch works with eager mode only\n            if not tf.executing_eagerly():\n                raise RuntimeError(\"tape.watch only works with eager mode\")\n            self.tape.watch(batch_embedding_tensor)\n            self.bet_ids_pair.append((batch_embedding_tensor, unique_ids))\n        outputs = tf.gather(batch_embedding_tensor, idx)\n        outputs = tf.reshape(\n            outputs, ids.get_shape().concatenate(self.output_dim)\n        )\n        return outputs\n\n    def reset(self):\n        self.bet_ids_pair = []\n        self.tape = None\n\n    def set_tape(self, tape):\n        self.tape = tape\n\n    def set_worker(self, worker):\n        self.worker = worker\n", "path": "elasticdl/python/elasticdl/layers/embedding.py"}]}