Repository: sql-machine-learning/elasticdl

Issue:
Prune invalid IDs (<0) before executing lookup_embedding.
We should keep the `elasticdl.layers.embedding` same as `tf.nn.safe_embedding_lookup_sparse`. https://www.tensorflow.org/api_docs/python/tf/nn/safe_embedding_lookup_sparse?version=stable


Discussion:

Resolving patch:
diff --git a/elasticdl/python/elasticdl/embedding_delegate.py b/elasticdl/python/elasticdl/embedding_delegate.py
--- a/elasticdl/python/elasticdl/embedding_delegate.py
+++ b/elasticdl/python/elasticdl/embedding_delegate.py
@@ -2,6 +2,8 @@
 
 import numpy as np
 import tensorflow as tf
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops, math_ops, sparse_ops
 
 EmbeddingAndIds = collections.namedtuple(
     "EmbeddingAndIds", ["batch_embedding", "batch_ids"]
@@ -27,9 +29,17 @@ def __init__(self, input_dim, output_dim, name):
         # different iterations.
         # `tf.Variable` requires initial value if shape has `None` dimension.
         self._embedding_and_ids_graph = []
+        self.tape = None
 
-    def init_for_graph_mode_if_necessary(self):
-        if tf.executing_eagerly() or self._embedding_and_ids_graph:
+    def set_tape(self, tape):
+        self.tape = tape
+
+    def _init_for_graph_mode_if_necessary(self):
+        if (
+            tf.executing_eagerly()
+            or self._embedding_and_ids_graph
+            or not self.tape
+        ):
             return
 
         self._embedding_and_ids_graph = [
@@ -51,7 +61,181 @@ def init_for_graph_mode_if_necessary(self):
             )
         ]
 
-    def lookup_embedding(self, unique_ids):
+    def embedding_lookup(self, ids):
+        """Looks up `ids` in a list of embedding tensors. The result of this
+        function is the same as Tensorflow `embedding_ops.embedding_lookup`.
+        But, this function is implemented to support lookup embedding using
+        ParameterServer distribution strategy.
+        """
+        self._init_for_graph_mode_if_necessary()
+
+        ids = tf.cast(ids, tf.int64)
+        ids = tf.convert_to_tensor(ids, name=self.name + "_ids")
+        flat_ids = tf.reshape(ids, [-1])
+        unique_ids, idx = tf.unique(flat_ids)
+        batch_embedding = self._get_embeddings_by_id(unique_ids)
+
+        # Gradient for `batch_embedding` is SparseTensor here due to
+        # `tf.gather` op. `tf.gather` accesses tensor slices, resulting in
+        # sparse tensor gradient.
+        # TODO: use tf.cond rather than python if statement
+        if self.tape:
+            batch_embedding = self._record_gradients(batch_embedding, flat_ids)
+
+        result = tf.gather(batch_embedding, idx)
+        # tf.reshape does not support shape with None. Replace None with -1.
+        if ids.get_shape().rank == 2:
+            input_length = ids.get_shape()[1]
+            if input_length is None:
+                result.set_shape(shape=(None, None, self.output_dim))
+                return result
+            output_shape = (-1, input_length, self.output_dim)
+        else:
+            output_shape = ids.get_shape().concatenate(self.output_dim)
+        result = tf.reshape(result, output_shape)
+        return result
+
+    def safe_embedding_lookup_sparse(
+        self, sparse_ids, sparse_weights=None, combiner="mean", default_id=None
+    ):
+        """Lookup embedding results, accounting for invalid IDs and empty
+        features. The result of this function is the same as
+        tf.nn.safe_embeddding_lookup_sparse`. But, this function is implemented
+        to support lookup embedding using ParameterServer distribution
+        strategy.
+        """
+        self._init_for_graph_mode_if_necessary()
+
+        sparse_ids = _prune_invalid_ids(sparse_ids)
+        # Fill in dummy values for empty features, if necessary.
+        sparse_ids, is_row_empty = sparse_ops.sparse_fill_empty_rows(
+            sparse_ids, 0
+        )
+        unique_ids, idx = tf.unique(sparse_ids.values)
+
+        segment_ids = sparse_ids.indices[:, 0]
+        if segment_ids.dtype != tf.int32:
+            segment_ids = tf.cast(segment_ids, tf.int32)
+
+        ids = sparse_ids.values
+        unique_ids, idx = tf.unique(ids)
+        batch_embedding = self._get_embeddings_by_id(unique_ids)
+
+        if sparse_weights is not None:
+            if self.tape:
+                batch_embedding = self._record_gradients(
+                    batch_embedding=batch_embedding, ids=ids
+                )
+
+            weights = sparse_weights.values
+            if weights.dtype != batch_embedding.dtype:
+                weights = math_ops.cast(weights, batch_embedding.dtype)
+
+            batch_embedding = array_ops.gather(batch_embedding, idx)
+            # Reshape weights to allow broadcast
+            ones = array_ops.fill(
+                array_ops.expand_dims(array_ops.rank(batch_embedding) - 1, 0),
+                1,
+            )
+            bcast_weights_shape = array_ops.concat(
+                [array_ops.shape(weights), ones], 0
+            )
+
+            orig_weights_shape = weights.get_shape()
+            weights = array_ops.reshape(weights, bcast_weights_shape)
+
+            # Set the weight shape, since after reshaping to
+            # bcast_weights_shape, the shape becomes None.
+            if batch_embedding.get_shape().ndims is not None:
+                weights.set_shape(
+                    orig_weights_shape.concatenate(
+                        [
+                            1
+                            for _ in range(
+                                batch_embedding.get_shape().ndims - 1
+                            )
+                        ]
+                    )
+                )
+
+            batch_embedding *= weights
+
+            if combiner == "sum":
+                batch_embedding = math_ops.segment_sum(
+                    batch_embedding, segment_ids
+                )
+            elif combiner == "mean":
+                batch_embedding = math_ops.segment_sum(
+                    batch_embedding, segment_ids
+                )
+                weight_sum = math_ops.segment_sum(weights, segment_ids)
+                batch_embedding = math_ops.div(batch_embedding, weight_sum)
+            elif combiner == "sqrtn":
+                batch_embedding = math_ops.segment_sum(
+                    batch_embedding, segment_ids
+                )
+                weights_squared = math_ops.pow(weights, 2)
+                weight_sum = math_ops.segment_sum(weights_squared, segment_ids)
+                weight_sum_sqrt = math_ops.sqrt(weight_sum)
+                batch_embedding = math_ops.div(
+                    batch_embedding, weight_sum_sqrt
+                )
+            else:
+                assert False, "Unrecognized combiner"
+        else:
+            if self.tape:
+                batch_embedding = self._record_gradients(
+                    batch_embedding=batch_embedding, ids=unique_ids,
+                )
+
+            assert idx is not None
+            if combiner == "sum":
+                batch_embedding = math_ops.sparse_segment_sum(
+                    batch_embedding, idx, segment_ids
+                )
+            elif combiner == "mean":
+                batch_embedding = math_ops.sparse_segment_mean(
+                    batch_embedding, idx, segment_ids
+                )
+            elif combiner == "sqrtn":
+                batch_embedding = math_ops.sparse_segment_sqrt_n(
+                    batch_embedding, idx, segment_ids
+                )
+            else:
+                assert False, "Unrecognized combiner"
+
+        # Broadcast is_row_empty to the same shape as embedding_lookup_result,
+        # for use in Select.
+        is_row_empty = array_ops.tile(
+            array_ops.reshape(is_row_empty, [-1, 1]),
+            array_ops.stack([1, array_ops.shape(batch_embedding)[1]]),
+        )
+
+        batch_embedding = array_ops.where(
+            is_row_empty,
+            array_ops.zeros_like(batch_embedding),
+            batch_embedding,
+            name=self.name,
+        )
+        batch_embedding.set_shape((None, self.output_dim))
+        return batch_embedding
+
+    def _get_embeddings_by_id(self, unique_ids):
+        """There is a memory leak when using tf.py_function with eager
+        execution. So, we only use tf.py_function in graph mode.
+        """
+        if isinstance(unique_ids, ops.EagerTensor):
+            batch_embedding = self._gather_embedding_vectors(unique_ids)
+            batch_embedding = tf.constant(batch_embedding, dtype=tf.float32)
+        else:
+            batch_embedding = tf.py_function(
+                self._gather_embedding_vectors,
+                inp=[unique_ids],
+                Tout=tf.float32,
+            )
+        return batch_embedding
+
+    def _gather_embedding_vectors(self, unique_ids):
         ids = unique_ids.numpy()
         self._check_id_valid(ids)
         if self._lookup_embedding_func:
@@ -70,9 +254,9 @@ def _check_id_valid(self, ids):
                 % (first_may_exceed_id, self.input_dim)
             )
 
-    def record_gradients(self, tape, batch_embedding, ids):
+    def _record_gradients(self, batch_embedding, ids):
         if tf.executing_eagerly():
-            tape.watch(batch_embedding)
+            self.tape.watch(batch_embedding)
             self._embedding_and_ids_eagerly.append(
                 EmbeddingAndIds(batch_embedding, ids)
             )
@@ -88,6 +272,7 @@ def record_gradients(self, tape, batch_embedding, ids):
         return batch_embedding
 
     def reset(self):
+        self.tape = None
         self._embedding_and_ids_eagerly = []
 
     def set_lookup_embedding_func(self, lookup_embedding_func):
@@ -101,3 +286,10 @@ def embedding_and_ids(self):
         if self._embedding_and_ids_eagerly:
             return self._embedding_and_ids_eagerly
         return self._embedding_and_ids_graph
+
+
+def _prune_invalid_ids(sparse_ids):
+    """Prune invalid IDs (< 0) from the input ids."""
+    is_id_valid = tf.greater_equal(sparse_ids.values, 0)
+    sparse_ids = sparse_ops.sparse_retain(sparse_ids, is_id_valid)
+    return sparse_ids
diff --git a/elasticdl/python/elasticdl/feature_column/feature_column.py b/elasticdl/python/elasticdl/feature_column/feature_column.py
--- a/elasticdl/python/elasticdl/feature_column/feature_column.py
+++ b/elasticdl/python/elasticdl/feature_column/feature_column.py
@@ -1,11 +1,10 @@
 import collections
 import math
 
-import tensorflow as tf
 from tensorflow.python.feature_column import feature_column as fc_old
 from tensorflow.python.feature_column import feature_column_v2 as fc_lib
 from tensorflow.python.framework import tensor_shape
-from tensorflow.python.ops import array_ops, init_ops, math_ops
+from tensorflow.python.ops import init_ops
 
 from elasticdl.python.elasticdl.embedding_delegate import EmbeddingDelegate
 
@@ -96,8 +95,6 @@ class EmbeddingColumn(
     ),
 ):
     def __init__(self, **kwargs):
-        self.tape = None
-
         default_num_buckets = (
             self.categorical_column.num_buckets
             if self._is_v2_column
@@ -152,9 +149,6 @@ def get_dense_tensor(self, transformation_cache, state_manager):
                 )
             )
 
-        if self.tape:
-            self._embedding_delegate.init_for_graph_mode_if_necessary()
-
         # Get sparse IDs and weights.
         sparse_tensors = self.categorical_column.get_sparse_tensors(
             transformation_cache, state_manager
@@ -163,111 +157,16 @@ def get_dense_tensor(self, transformation_cache, state_manager):
         # Look up the embedding from the sparse input
         sparse_ids = sparse_tensors.id_tensor
         sparse_weights = sparse_tensors.weight_tensor
-
-        segment_ids = sparse_ids.indices[:, 0]
-        if segment_ids.dtype != tf.int32:
-            segment_ids = tf.cast(segment_ids, tf.int32)
-
-        ids = sparse_ids.values
-        unique_ids, idx = tf.unique(ids)
-
-        batch_embedding = tf.py_function(
-            self.lookup_embedding, inp=[unique_ids], Tout=tf.float32
+        result = self._embedding_delegate.safe_embedding_lookup_sparse(
+            sparse_ids, sparse_weights=sparse_weights, combiner=self.combiner
         )
-
-        if sparse_weights is not None:
-            if self.tape:
-                batch_embedding = self._embedding_delegate.record_gradients(
-                    tape=self.tape, batch_embedding=batch_embedding, ids=ids
-                )
-
-            weights = sparse_weights.values
-            if weights.dtype != batch_embedding.dtype:
-                weights = math_ops.cast(weights, batch_embedding.dtype)
-
-            batch_embedding = array_ops.gather(batch_embedding, idx)
-
-            # Reshape weights to allow broadcast
-            ones = array_ops.fill(
-                array_ops.expand_dims(array_ops.rank(batch_embedding) - 1, 0),
-                1,
-            )
-            bcast_weights_shape = array_ops.concat(
-                [array_ops.shape(weights), ones], 0
-            )
-
-            orig_weights_shape = weights.get_shape()
-            weights = array_ops.reshape(weights, bcast_weights_shape)
-
-            # Set the weight shape, since after reshaping to
-            # bcast_weights_shape, the shape becomes None.
-            if batch_embedding.get_shape().ndims is not None:
-                weights.set_shape(
-                    orig_weights_shape.concatenate(
-                        [
-                            1
-                            for _ in range(
-                                batch_embedding.get_shape().ndims - 1
-                            )
-                        ]
-                    )
-                )
-
-            batch_embedding *= weights
-
-            if self.combiner == "sum":
-                batch_embedding = math_ops.segment_sum(
-                    batch_embedding, segment_ids
-                )
-            elif self.combiner == "mean":
-                batch_embedding = math_ops.segment_sum(
-                    batch_embedding, segment_ids
-                )
-                weight_sum = math_ops.segment_sum(weights, segment_ids)
-                batch_embedding = math_ops.div(batch_embedding, weight_sum)
-            elif self.combiner == "sqrtn":
-                batch_embedding = math_ops.segment_sum(
-                    batch_embedding, segment_ids
-                )
-                weights_squared = math_ops.pow(weights, 2)
-                weight_sum = math_ops.segment_sum(weights_squared, segment_ids)
-                weight_sum_sqrt = math_ops.sqrt(weight_sum)
-                batch_embedding = math_ops.div(
-                    batch_embedding, weight_sum_sqrt
-                )
-            else:
-                assert False, "Unrecognized combiner"
-        else:
-            if self.tape:
-                batch_embedding = self._embedding_delegate.record_gradients(
-                    tape=self.tape,
-                    batch_embedding=batch_embedding,
-                    ids=unique_ids,
-                )
-
-            assert idx is not None
-            if self.combiner == "sum":
-                batch_embedding = math_ops.sparse_segment_sum(
-                    batch_embedding, idx, segment_ids
-                )
-            elif self.combiner == "mean":
-                batch_embedding = math_ops.sparse_segment_mean(
-                    batch_embedding, idx, segment_ids
-                )
-            elif self.combiner == "sqrtn":
-                batch_embedding = math_ops.sparse_segment_sqrt_n(
-                    batch_embedding, idx, segment_ids
-                )
-            else:
-                assert False, "Unrecognized combiner"
-
-        return batch_embedding
+        return result
 
     def lookup_embedding(self, unique_ids):
         return self._embedding_delegate.lookup_embedding(unique_ids)
 
     def set_tape(self, tape):
-        self.tape = tape
+        self._embedding_delegate.set_tape(tape)
 
     def set_lookup_embedding_func(self, func):
         """Sets function for looking up embeddings in the PS.
@@ -281,7 +180,6 @@ def set_lookup_embedding_func(self, func):
         self._embedding_delegate.set_lookup_embedding_func(func)
 
     def reset(self):
-        self.tape = None
         self._embedding_delegate.reset()
 
     @property
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
@@ -1,9 +1,7 @@
-import numpy as np
 import tensorflow as tf
-from tensorflow.python.framework import ops
 from tensorflow.python.keras.utils import tf_utils
 
-from elasticdl.python.elasticdl.embedding_delegate import EmbeddingAndIds
+from elasticdl.python.elasticdl.embedding_delegate import EmbeddingDelegate
 
 
 class Embedding(tf.keras.layers.Layer):
@@ -63,26 +61,9 @@ def __init__(
         # different iterations.
         # `tf.Variable` requires initial value if shape has `None` dimension.
         self._embedding_and_ids_graph = []
-
-    def _init_for_graph_mode(self):
-        self._embedding_and_ids_graph = [
-            EmbeddingAndIds(
-                batch_embedding=tf.Variable(
-                    # In some cases, `tf.Variable` requires that initial value
-                    # is callable.
-                    initial_value=lambda: tf.zeros((1, self.output_dim)),
-                    shape=tf.TensorShape((None, self.output_dim)),
-                    dtype=tf.float32,
-                    trainable=True,
-                ),
-                batch_ids=tf.Variable(
-                    initial_value=lambda: tf.zeros((1, 1), dtype=tf.int64),
-                    shape=tf.TensorShape(None),
-                    dtype=tf.int64,
-                    trainable=False,
-                ),
-            )
-        ]
+        self.embedding_delegate = EmbeddingDelegate(
+            self.input_dim, self.output_dim, self.name
+        )
 
     @tf_utils.shape_type_conversion
     def compute_output_shape(self, input_shape):
@@ -122,125 +103,21 @@ def name(self):
     def get_key(name_list):
         return "-".join(map(str, name_list))
 
-    def lookup_embedding(self, unique_ids):
-        ids = unique_ids.numpy()
-        self._check_id_valid(ids)
-        if self._lookup_embedding_func:
-            embedding_vectors = self._lookup_embedding_func(self._name, ids)
-            return embedding_vectors
-
-    def _check_id_valid(self, ids):
-        if not self.input_dim:
-            return
-        first_may_exceed_id = ids[np.argmax(ids >= self.input_dim)]
-        if self.input_dim is not None and first_may_exceed_id > self.input_dim:
-            raise ValueError(
-                " The embedding id cannot be bigger "
-                "than input_dim. id = %d is not in [0, %d)"
-                % (first_may_exceed_id, self.input_dim)
-            )
-        if any(ids < 0):
-            raise ValueError(
-                "The embedding id cannot be less than zero. "
-                "Found invalid id %d" % ids[ids < 0][0]
-            )
-
-    def _record_gradients(self, batch_embedding, ids):
-        if tf.executing_eagerly():
-            self.tape.watch(batch_embedding)
-            self._embedding_and_ids_eagerly.append(
-                EmbeddingAndIds(batch_embedding, ids)
-            )
-        else:
-            # In graph mode, assigning tensors to trainable variables is
-            # allowed and tape can record the gradients of trainable
-            # variables automatically.
-            embedding_and_ids = self._embedding_and_ids_graph[0]
-            embedding_and_ids.batch_embedding.assign(batch_embedding)
-            embedding_and_ids.batch_ids.assign(ids)
-            batch_embedding = embedding_and_ids.batch_embedding
-        return batch_embedding
-
     def call(self, input):
-        if (
-            self.tape
-            and not tf.executing_eagerly()
-            and not self._embedding_and_ids_graph
-        ):
-            self._init_for_graph_mode()
-
         input = tf.cast(input, tf.int64)
         if isinstance(input, tf.SparseTensor):
             return self._sparse_input_call(input)
-
-        ids = tf.convert_to_tensor(input, name="embedding_ids")
-        flat_ids = tf.reshape(ids, [-1])
-        unique_ids, idx = tf.unique(flat_ids)
-
-        # There is a memory leak when using tf.py_function with eager
-        # execution. So, we only use tf.py_function in graph mode.
-        if isinstance(unique_ids, ops.EagerTensor):
-            batch_embedding = self.lookup_embedding(unique_ids)
-            batch_embedding = tf.constant(batch_embedding)
         else:
-            batch_embedding = tf.py_function(
-                self.lookup_embedding, inp=[unique_ids], Tout=tf.float32
-            )
-
-        # Gradient for `batch_embedding` is SparseTensor here due to
-        # `tf.gather` op. `tf.gather` accesses tensor slices, resulting in
-        # sparse tensor gradient.
-        # TODO: use tf.cond rather than python if statement
-        if self.tape:
-            batch_embedding = self._record_gradients(batch_embedding, flat_ids)
-
-        outputs = tf.gather(batch_embedding, idx)
-        # tf.reshape does not support shape with None. Replace None with -1.
-        if ids.get_shape().rank == 2:
-            input_length = ids.get_shape()[1]
-            if input_length is None:
-                outputs.set_shape(shape=(None, None, self.output_dim))
-                return outputs
-            output_shape = (-1, input_length, self.output_dim)
-        else:
-            output_shape = ids.get_shape().concatenate(self.output_dim)
-        outputs = tf.reshape(outputs, output_shape)
-        # TODO: support combiner for dense input
-        return outputs
+            return self.embedding_delegate.embedding_lookup(input)
 
     def _sparse_input_call(self, sparse_input):
         if self.combiner not in ["sum", "mean", "sqrtn"]:
             raise ValueError(
                 "combiner must set sum, mean or sqrtn for sparse input"
             )
-        unique_ids, idx = tf.unique(sparse_input.values)
-        # Gradient for `batch_embedding` is dense tensor.
-        batch_embedding = tf.py_function(
-            self.lookup_embedding, inp=[unique_ids], Tout=tf.float32
+        batch_embedding = self.embedding_delegate.safe_embedding_lookup_sparse(
+            sparse_input, combiner=self.combiner
         )
-        # TODO: use tf.cond rather than python if statement
-        if self.tape:
-            batch_embedding = self._record_gradients(
-                batch_embedding, unique_ids
-            )
-
-        segment_ids = sparse_input.indices[:, 0]
-        if segment_ids.dtype != tf.int32:
-            segment_ids = tf.cast(segment_ids, tf.int32)
-
-        if self.combiner == "sum":
-            batch_embedding = tf.sparse.segment_sum(
-                batch_embedding, idx, segment_ids
-            )
-        elif self.combiner == "mean":
-            batch_embedding = tf.sparse.segment_mean(
-                batch_embedding, idx, segment_ids
-            )
-        elif self.combiner == "sqrtn":
-            batch_embedding = tf.sparse.segment_sqrt_n(
-                batch_embedding, idx, segment_ids
-            )
-        batch_embedding.set_shape((None, self.output_dim))
         return batch_embedding
 
     def compute_mask(self, inputs, mask=None):
@@ -251,28 +128,21 @@ def compute_mask(self, inputs, mask=None):
         return tf.math.not_equal(inputs, 0)
 
     def reset(self):
-        self._embedding_and_ids_eagerly = []
-        self.tape = None
+        self.embedding_delegate.reset()
 
     def set_tape(self, tape):
-        self.tape = tape
+        self.embedding_delegate.set_tape(tape)
 
     def set_lookup_embedding_func(self, func):
         """Sets function for looking up embeddings in the PS.
-
         Args:
             func: The function used to look up embeddings. The arguments of
                 are `(layer_name, embedding_id_list)`, where `layer_name` is
                 the name of embedding layer, and `embedding_id_list` is a list
                 of embedding ids to be looked up.
         """
-        self._lookup_embedding_func = func
+        self.embedding_delegate.set_lookup_embedding_func(func)
 
     @property
     def embedding_and_ids(self):
-        """
-        Return bet and ids pairs.
-        """
-        if self._embedding_and_ids_eagerly:
-            return self._embedding_and_ids_eagerly
-        return self._embedding_and_ids_graph
+        return self.embedding_delegate.embedding_and_ids


Test patch:
diff --git a/elasticdl/python/tests/feature_column_test.py b/elasticdl/python/tests/feature_column_test.py
--- a/elasticdl/python/tests/feature_column_test.py
+++ b/elasticdl/python/tests/feature_column_test.py
@@ -40,11 +40,11 @@ def test_call_embedding_column(self):
             ),
             dimension=dimension,
         )
-        item_id_embedding.lookup_embedding = lambda unique_ids: (
-            generate_vectors_with_one_hot_value(
-                unique_ids, item_id_embedding.dimension
-            )
-        )
+
+        def _mock_gather_embedding(name, ids):
+            return generate_vectors_with_one_hot_value(ids, dimension)
+
+        item_id_embedding.set_lookup_embedding_func(_mock_gather_embedding)
 
         output = call_feature_columns(
             [item_id_embedding], {"item_id": [1, 2, 3]}
@@ -71,11 +71,11 @@ def test_call_embedding_column_with_weights(self):
             initializer=tf.initializers.identity,
             combiner="sum",
         )
-        item_id_embedding.lookup_embedding = lambda unique_ids: (
-            generate_vectors_with_one_hot_value(
-                unique_ids, item_id_embedding.dimension
-            )
-        )
+
+        def _mock_gather_embedding(name, ids):
+            return generate_vectors_with_one_hot_value(ids, dimension)
+
+        item_id_embedding.set_lookup_embedding_func(_mock_gather_embedding)
 
         output = call_feature_columns(
             [item_id_embedding],
@@ -125,11 +125,11 @@ def test_embedding_column_gradients(self):
                 dimension=dimension,
                 combiner=combiner,
             )
-            item_id_embedding.lookup_embedding = lambda unique_ids: (
-                generate_vectors_fill_with_id_value(
-                    unique_ids, item_id_embedding.dimension
-                )
-            )
+
+            def _mock_gather_embedding(name, ids):
+                return generate_vectors_fill_with_id_value(ids, dimension)
+
+            item_id_embedding.set_lookup_embedding_func(_mock_gather_embedding)
 
             dense_features = tf.keras.layers.DenseFeatures([item_id_embedding])
             call_fns = [dense_features.call, tf.function(dense_features.call)]
@@ -221,11 +221,11 @@ def test_embedding_column_gradients_with_weights(self):
                 initializer=tf.initializers.identity,
                 combiner=combiner,
             )
-            item_id_embedding.lookup_embedding = lambda unique_ids: (
-                generate_vectors_fill_with_id_value(
-                    unique_ids, item_id_embedding.dimension
-                )
-            )
+
+            def _mock_gather_embedding(name, ids):
+                return generate_vectors_fill_with_id_value(ids, dimension)
+
+            item_id_embedding.set_lookup_embedding_func(_mock_gather_embedding)
 
             dense_features = tf.keras.layers.DenseFeatures([item_id_embedding])
             call_fns = [dense_features.call, tf.function(dense_features.call)]
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
@@ -65,7 +65,7 @@ def __init__(self, embedding_size, output_dim):
         for i in range(embedding_size):
             self.embedding[i].fill(i)
 
-    def lookup_embedding(self, ids):
+    def lookup_embedding(self, name, ids):
         values = np.take(self.embedding, ids, axis=0)
         return values
 
@@ -84,7 +84,7 @@ def create_embedding_layer(
         mask_zero=mask_zero,
     )
     worker = mock_worker(embedding_size, output_dim)
-    layer.lookup_embedding = worker.lookup_embedding
+    layer.set_lookup_embedding_func(worker.lookup_embedding)
     return layer