prompt: You will be provided with a partial code base and an issue statement explaining a problem to resolve.
  <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,
  )
  ```
  
  </issue>
  <code>
  [start of elasticdl/python/elasticdl/layers/embedding.py]
  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 of elasticdl/python/elasticdl/layers/embedding.py]
  </code>
  
  I need you to solve this issue by generating a single patch file that I can apply directly to this repository using git apply. Please respond with a single patch file in the following format.
  <patch>
  diff --git a/file.py b/file.py
  --- a/file.py
  +++ b/file.py
  @@ -1,27 +1,35 @@
   def euclidean(a, b):
  -    while b:
  -        a, b = b, a % b
  -    return a
  +    if b == 0:
  +        return a
  +    return euclidean(b, a % b)
   
   
   def bresenham(x0, y0, x1, y1):
       points = []
       dx = abs(x1 - x0)
       dy = abs(y1 - y0)
  -    sx = 1 if x0 < x1 else -1
  -    sy = 1 if y0 < y1 else -1
  -    err = dx - dy
  +    x, y = x0, y0
  +    sx = -1 if x0 > x1 else 1
  +    sy = -1 if y0 > y1 else 1
   
  -    while True:
  -        points.append((x0, y0))
  -        if x0 == x1 and y0 == y1:
  -            break
  -        e2 = 2 * err
  -        if e2 > -dy:
  -            err -= dy
  -            x0 += sx
  -        if e2 < dx:
  -            err += dx
  -            y0 += sy
  +    if dx > dy:
  +        err = dx / 2.0
  +        while x != x1:
  +            points.append((x, y))
  +            err -= dy
  +            if err < 0:
  +                y += sy
  +                err += dx
  +            x += sx
  +    else:
  +        err = dy / 2.0
  +        while y != y1:
  +            points.append((x, y))
  +            err -= dx
  +            if err < 0:
  +                x += sx
  +                err += dy
  +            y += sy
  +
  +    points.append((x, y))
       return points
  
  </patch>
golden_diff: 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
  
verification_info: {"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"}]}