Repository: sql-machine-learning/elasticdl

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,
)
```


Discussion:

Resolving patch:
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 @@ class Embedding(tf.keras.layers.Layer):
     """
     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 @@ def call(self, input):
             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
 


Test patch:
diff --git a/elasticdl/python/tests/embedding_test_module.py b/elasticdl/python/tests/embedding_test_module.py
--- a/elasticdl/python/tests/embedding_test_module.py
+++ b/elasticdl/python/tests/embedding_test_module.py
@@ -5,13 +5,13 @@
 
 
 class CustomModel(tf.keras.Model):
-    def __init__(self, embedding_dim=16):
+    def __init__(self, output_dim=16):
         super(CustomModel, self).__init__(name="embedding_test_model")
-        self.embedding_dim = embedding_dim
-        self.embedding_1 = Embedding(embedding_dim)
-        self.embedding_2 = Embedding(embedding_dim)
+        self.output_dim = output_dim
+        self.embedding_1 = Embedding(output_dim)
+        self.embedding_2 = Embedding(output_dim)
         self.concat = Concatenate()
-        self.dense = Dense(1, input_shape=(embedding_dim * 3,))
+        self.dense = Dense(1, input_shape=(output_dim * 3,))
         self.flatten = Flatten()
 
     def call(self, inputs, training=False):
diff --git a/elasticdl/python/tests/layer_test.py b/elasticdl/python/tests/layer_test.py
--- a/elasticdl/python/tests/layer_test.py
+++ b/elasticdl/python/tests/layer_test.py
@@ -56,11 +56,11 @@ def test_find_layer_nested(self):
 
 
 class mock_worker:
-    def __init__(self, embedding_size, embedding_dim):
+    def __init__(self, embedding_size, output_dim):
         self.embedding_size = embedding_size
-        self.embedding_dim = embedding_dim
+        self.output_dim = output_dim
         self.embedding = np.ndarray(
-            shape=(embedding_size, embedding_dim), dtype=np.float32
+            shape=(embedding_size, output_dim), dtype=np.float32
         )
         for i in range(embedding_size):
             self.embedding[i].fill(i)
@@ -70,9 +70,9 @@ def embedding_lookup(self, ids, name, embedding_initializer):
         return values
 
 
-def create_embedding_layer(embedding_size, embedding_dim):
-    layer = Embedding(embedding_dim)
-    worker = mock_worker(embedding_size, embedding_dim)
+def create_embedding_layer(embedding_size, output_dim, input_length=None):
+    layer = Embedding(output_dim, input_length=input_length)
+    worker = mock_worker(embedding_size, output_dim)
     layer.set_worker(worker)
     return layer
 
@@ -84,19 +84,19 @@ def layer_call(layer, inputs):
 
 class EmbeddingLayerTest(unittest.TestCase):
     def test_embedding_layer(self):
-        embedding_dim = 8
+        output_dim = 8
         embedding_size = 16
-        layer = create_embedding_layer(embedding_size, embedding_dim)
+        layer = create_embedding_layer(embedding_size, output_dim)
 
         input_shape = (2, 6)
         output_shape = layer.compute_output_shape(input_shape)
-        self.assertEqual(output_shape, input_shape + (embedding_dim,))
+        self.assertEqual(output_shape, input_shape + (output_dim,))
 
         ids = [0, 1, 3, 8, 3, 2, 3]
         values = layer.call(ids)
         values = values.numpy()
         for index, idx in enumerate(ids):
-            correct_value = np.array([idx] * embedding_dim, dtype=np.float32)
+            correct_value = np.array([idx] * output_dim, dtype=np.float32)
             self.assertTrue((values[index] == correct_value).all())
 
         results = layer_call(layer, ids)
@@ -108,22 +108,35 @@ def test_embedding_layer(self):
         outputs = model.call(np.array([ids, ids]))
         values = outputs.numpy()[1]
         for index, idx in enumerate(ids):
-            correct_value = np.array([idx] * embedding_dim, dtype=np.float32)
+            correct_value = np.array([idx] * output_dim, dtype=np.float32)
+            self.assertTrue((values[index] == correct_value).all())
+
+    def test_embedding_layer_with_input_length(self):
+        output_dim = 8
+        embedding_size = 16
+        input_length = 4
+        layer = create_embedding_layer(
+            embedding_size, output_dim, input_length=input_length
+        )
+        ids = [[0, 1, 3, 8], [5, 3, 2, 3]]
+        flatten_ids = ids[0] + ids[1]
+        values = layer.call(ids)
+        values = values.numpy().reshape(-1, output_dim)
+        for index, idx in enumerate(flatten_ids):
+            correct_value = np.array([idx] * output_dim, dtype=np.float32)
             self.assertTrue((values[index] == correct_value).all())
 
     def test_embedding_layer_gradient(self):
-        embedding_dim = 8
+        output_dim = 8
         embedding_size = 16
-        layer = create_embedding_layer(embedding_size, embedding_dim)
+        layer = create_embedding_layer(embedding_size, output_dim)
         inputs_list = [
             tf.keras.backend.constant([[0, 1, 3], [1, 2, 0]], dtype=tf.int64),
             tf.keras.backend.constant(
                 [[0, 10, 3], [11, 2, 1]], dtype=tf.int64
             ),
         ]
-        multiply_values = np.ndarray(
-            shape=(6, embedding_dim), dtype=np.float32
-        )
+        multiply_values = np.ndarray(shape=(6, output_dim), dtype=np.float32)
         for i in range(6):
             multiply_values[i].fill(i)
         multiply_tensor = tf.convert_to_tensor(multiply_values)