Repository: sql-machine-learning/elasticdl

Issue:
Support mask_zero argument in ElasticDL embedding layer
Support `mask_zero` argument in `elasticdl.Embedding`.

`mask_zero` is an argument in keras Embedding layer and it is used for inputs with different `input_length` in one minibatch. More details can be found in [keras Embedding doc](https://keras.io/layers/embeddings/).


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
@@ -41,7 +41,6 @@ def __init__(
 
         self.output_dim = output_dim
         self.embedding_initializer = embedding_initializer
-        # TODO: support mask_zero
         self.supports_masking = mask_zero
         self.input_length = input_length
         self.combiner = combiner
@@ -147,6 +146,13 @@ def _sparse_input_call(self, sparse_input):
             embeddings = tf.sparse.segment_sqrt_n(embeddings, idx, segment_ids)
         return embeddings
 
+    def compute_mask(self, inputs, mask=None):
+        if isinstance(input, tf.SparseTensor):
+            raise ValueError("SparseTensor inputs do not support mask_zero")
+        if not self.supports_masking:
+            return None
+        return tf.math.not_equal(inputs, 0)
+
     def reset(self):
         self.bet_ids_pair = []
         self.tape = None


Test patch:
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
@@ -71,9 +71,18 @@ def embedding_lookup(self, ids, name, embedding_initializer):
 
 
 def create_embedding_layer(
-    embedding_size, output_dim, input_length=None, combiner=None
+    embedding_size,
+    output_dim,
+    input_length=None,
+    combiner=None,
+    mask_zero=False,
 ):
-    layer = Embedding(output_dim, input_length=input_length, combiner=combiner)
+    layer = Embedding(
+        output_dim,
+        input_length=input_length,
+        combiner=combiner,
+        mask_zero=mask_zero,
+    )
     worker = mock_worker(embedding_size, output_dim)
     layer.set_worker(worker)
     return layer
@@ -173,6 +182,30 @@ def test_embedding_layer_with_sparse_input(self):
             for n, v in enumerate(correct_values):
                 self.assertAlmostEqual(outputs[n][0], v, place)
 
+    def test_embedding_layer_with_mask_zero(self):
+        output_dim = 8
+        embedding_size = 16
+        mask_zero = True
+        layer = create_embedding_layer(
+            embedding_size, output_dim, mask_zero=mask_zero
+        )
+        ids = [[0, 1, 3, 8], [5, 0, 2, 3]]
+        correct_masks = np.ones((2, 4), dtype=np.bool)
+        correct_masks[0][0] = False
+        correct_masks[1][1] = False
+
+        masks = layer.compute_mask(ids)
+        self.assertTrue((masks.numpy() == correct_masks).all())
+
+        # SparseTensor inputs will raise error
+        indices = [[0, 0], [1, 1], [1, 2], [2, 1], [3, 1], [3, 3], [3, 5]]
+        values = [1, 1, 3, 2, 0, 2, 6]
+        dense_shape = [4, 6]
+        inputs = tf.SparseTensor(
+            indices=indices, values=values, dense_shape=dense_shape
+        )
+        self.assertRaises(ValueError, layer.compute_mask, inputs)
+
     def test_embedding_layer_gradient(self):
         output_dim = 8
         embedding_size = 16