We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Do not save ElasticDL embedding's variables in the master
In current code, ElasticDL embedding layer create trainable variables and use them only for training under graph mode (i.e. under tf.function decorator). These variables are used for training and thus should only be created in workers. 

However, now ElasticDL embedding layer creates these variables in `__init__` function. The master (or PS) creates these variables when it creates a model instance, and then save these variables in `self._model`. 

I propose that ElasticDL embedding layer should not create these variables in `__init__` function. Instead, creating them before the training process only in worker is enough.

--- 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 collections
2 
3 import numpy as np
4 import tensorflow as tf
5 from tensorflow.python.keras.utils import tf_utils
6 
7 from elasticdl.python.master.embedding_service import EmbeddingService
8 
9 EmbeddingAndIds = collections.namedtuple(
10     "EmbeddingAndIds", ["batch_embedding", "batch_ids"]
11 )
12 
13 
14 class Embedding(tf.keras.layers.Layer):
15     """
16     Input: indexes for the embedding entries with a shape of
17       (batch_size, input_length). Input can be either dense tensor
18       or SparseTensor.
19     Output:
20       corresponding (combined) embeddings with a shape of
21       (batch_size, input_length, output_dim) if combiner is None
22       (batch_size, output_dim) if combiner is not None
23     Arguments:
24       output_dim: the dimension of the embedding vector
25       embedding_initializer: Initializer for embedding table
26       mask_zero: Whether or not the input value 0 is a special "padding"
27         value that should be masked out.
28         If input is SparseTensor, mask_zero must be False.
29       input_length: Length of input sequences, when it is constant.
30         This argument is required if you are going to connect
31         `Flatten` then `Dense` layers upstream
32         (without it, the shape of the dense outputs cannot be computed).
33       combiner: A string specifying the reduction op or None if not used.
34         "mean", "sqrtn" and "sum" are supported for the reduction op.
35         If input is SparseTensor, combiner must set as a reduction op.
36     """
37 
38     def __init__(
39         self,
40         output_dim,
41         embedding_initializer="uniform",
42         mask_zero=False,
43         input_length=None,
44         combiner=None,
45         embedding_service_endpoint=None,
46         **kwargs
47     ):
48         if "input_shape" not in kwargs and input_length:
49             kwargs["input_shape"] = (input_length,)
50         super(Embedding, self).__init__(**kwargs)
51 
52         self.output_dim = output_dim
53         self.embedding_initializer = embedding_initializer
54         self.supports_masking = mask_zero
55         self.input_length = input_length
56         self.combiner = combiner
57         self.embedding_service_endpoint = embedding_service_endpoint
58         self.tape = None
59         self.lookup_func = None
60 
61         self._embedding_and_ids_eagerly = []
62 
63         # BET's shape and ids' shape in `self._embedding_and_ids_graph` have
64         # `None` dimension. This is because they have different shapes in
65         # different iterations.
66         # `tf.Variable` requires initial value if shape has `None` dimension.
67         self._embedding_and_ids_graph = [
68             EmbeddingAndIds(
69                 batch_embedding=tf.Variable(
70                     initial_value=tf.zeros((1, self.output_dim)),
71                     shape=tf.TensorShape((None, self.output_dim)),
72                     dtype=tf.float32,
73                     trainable=True,
74                 ),
75                 batch_ids=tf.Variable(
76                     initial_value=tf.zeros((1, 1), dtype=tf.int64),
77                     shape=tf.TensorShape(None),
78                     dtype=tf.int64,
79                     trainable=False,
80                 ),
81             )
82         ]
83 
84     @tf_utils.shape_type_conversion
85     def compute_output_shape(self, input_shape):
86         # this function is taken from
87         # tf.keras.layers.Embedding.compute_output_shape
88         # https://github.com/tensorflow/tensorflow/blob/3f3c728bf80e0fd6653744318cbbfe1454c6ddca/tensorflow/python/keras/layers/embeddings.py#L156
89         if self.input_length is None:
90             return input_shape + (self.output_dim,)
91         else:
92             if isinstance(self.input_length, (list, tuple)):
93                 in_lens = list(self.input_length)
94             else:
95                 in_lens = [self.input_length]
96             if len(in_lens) != len(input_shape) - 1:
97                 raise ValueError(
98                     '"input_length" is %s, '
99                     "but received input has shape %s"
100                     % (str(self.input_length), str(input_shape))
101                 )
102             else:
103                 for i, (s1, s2) in enumerate(zip(in_lens, input_shape[1:])):
104                     if s1 is not None and s2 is not None and s1 != s2:
105                         raise ValueError(
106                             '"input_length" is %s, '
107                             "but received input has shape %s"
108                             % (str(self.input_length), str(input_shape))
109                         )
110                     elif s1 is None:
111                         in_lens[i] = s2
112             return (input_shape[0],) + tuple(in_lens) + (self.output_dim,)
113 
114     @property
115     def name(self):
116         return self._name
117 
118     @staticmethod
119     def get_key(name_list):
120         return "-".join(map(str, name_list))
121 
122     def lookup_embedding(self, unique_ids):
123         ids = unique_ids.numpy()
124         keys = [Embedding.get_key([self._name, id]) for id in ids]
125         (
126             embedding_vectors,
127             unknown_keys_index,
128         ) = EmbeddingService.lookup_embedding(
129             keys=keys,
130             embedding_service_endpoint=self.embedding_service_endpoint,
131         )
132 
133         if unknown_keys_index:
134             # Initialize unknown_keys' embedding vectors and write into Redis.
135             unknown_keys = [keys[index] for index in unknown_keys_index]
136             initializer = tf.keras.initializers.get(self.embedding_initializer)
137             embedding_vector_init = [
138                 initializer(shape=[1, self.output_dim]).numpy()
139                 for _ in unknown_keys
140             ]
141             embedding_vector_init = np.concatenate(
142                 embedding_vector_init, axis=0
143             )
144             EmbeddingService.update_embedding(
145                 keys=unknown_keys,
146                 embedding_vectors=embedding_vector_init,
147                 embedding_service_endpoint=self.embedding_service_endpoint,
148                 set_if_not_exist=True,
149             )
150             # Lookup unknown_keys' embedding vectors
151             (
152                 embedding_vectors_new,
153                 unknown_keys_idx_new,
154             ) = EmbeddingService.lookup_embedding(
155                 keys=unknown_keys,
156                 embedding_service_endpoint=self.embedding_service_endpoint,
157             )
158             if unknown_keys_idx_new:
159                 raise Exception(
160                     "Update embedding vector: %s failed."
161                     % str(
162                         [unknown_keys[index] for index in unknown_keys_idx_new]
163                     )
164                 )
165             for key_index, vector in zip(
166                 unknown_keys_index, embedding_vectors_new
167             ):
168                 embedding_vectors[key_index] = vector
169         embedding_vectors = np.concatenate(embedding_vectors, axis=0)
170         return embedding_vectors.reshape((len(keys), self.output_dim))
171 
172     def _record_gradients(self, batch_embedding, ids):
173         if tf.executing_eagerly():
174             self.tape.watch(batch_embedding)
175             self._embedding_and_ids_eagerly.append(
176                 EmbeddingAndIds(batch_embedding, ids)
177             )
178         else:
179             # In graph mode, assigning tensors to trainable variables is
180             # allowed and tape can record the gradients of trainable
181             # variables automatically.
182             embedding_and_ids = self._embedding_and_ids_graph[0]
183             embedding_and_ids.batch_embedding.assign(batch_embedding)
184             embedding_and_ids.batch_ids.assign(ids)
185             batch_embedding = embedding_and_ids.batch_embedding
186         return batch_embedding
187 
188     def call(self, input):
189         input = tf.cast(input, tf.int64)
190         if isinstance(input, tf.SparseTensor):
191             return self._sparse_input_call(input)
192 
193         ids = tf.convert_to_tensor(input, name="embedding_ids")
194         flat_ids = tf.reshape(ids, [-1])
195         unique_ids, idx = tf.unique(flat_ids)
196         # Gradient for `batch_embedding` is SparseTensor here due to
197         # `tf.gather` op. `tf.gather` accesses tensor slices, resulting in
198         # sparse tensor gradient.
199         batch_embedding = tf.py_function(
200             self.lookup_embedding, inp=[unique_ids], Tout=tf.float32
201         )
202         # TODO: use tf.cond rather than python if statement
203         if self.tape:
204             batch_embedding = self._record_gradients(batch_embedding, flat_ids)
205 
206         outputs = tf.gather(batch_embedding, idx)
207         # tf.reshape does not support shape with None. Replace None with -1.
208         if ids.get_shape().rank == 2:
209             input_length = ids.get_shape()[1]
210             if input_length is None:
211                 outputs.set_shape(shape=(None, None, self.output_dim))
212                 return outputs
213             output_shape = (-1, input_length, self.output_dim)
214         else:
215             output_shape = ids.get_shape().concatenate(self.output_dim)
216         outputs = tf.reshape(outputs, output_shape)
217         # TODO: support combiner for dense input
218         return outputs
219 
220     def _sparse_input_call(self, sparse_input):
221         if self.combiner not in ["sum", "mean", "sqrtn"]:
222             raise ValueError(
223                 "combiner must set sum, mean or sqrtn for sparse input"
224             )
225         unique_ids, idx = tf.unique(sparse_input.values)
226         # Gradient for `batch_embedding` is dense tensor.
227         batch_embedding = tf.py_function(
228             self.lookup_embedding, inp=[unique_ids], Tout=tf.float32
229         )
230         # TODO: use tf.cond rather than python if statement
231         if self.tape:
232             batch_embedding = self._record_gradients(
233                 batch_embedding, unique_ids
234             )
235 
236         segment_ids = sparse_input.indices[:, 0]
237         if segment_ids.dtype != tf.int32:
238             segment_ids = tf.cast(segment_ids, tf.int32)
239 
240         if self.combiner == "sum":
241             batch_embedding = tf.sparse.segment_sum(
242                 batch_embedding, idx, segment_ids
243             )
244         elif self.combiner == "mean":
245             batch_embedding = tf.sparse.segment_mean(
246                 batch_embedding, idx, segment_ids
247             )
248         elif self.combiner == "sqrtn":
249             batch_embedding = tf.sparse.segment_sqrt_n(
250                 batch_embedding, idx, segment_ids
251             )
252         return batch_embedding
253 
254     def compute_mask(self, inputs, mask=None):
255         if isinstance(input, tf.SparseTensor):
256             raise ValueError("SparseTensor inputs do not support mask_zero")
257         if not self.supports_masking:
258             return None
259         return tf.math.not_equal(inputs, 0)
260 
261     def reset(self):
262         self._embedding_and_ids_eagerly = []
263         self.tape = None
264 
265     def set_tape(self, tape):
266         self.tape = tape
267 
268     def set_endpoint(self, endpoint):
269         self.embedding_service_endpoint = endpoint
270 
271     @property
272     def embedding_and_ids(self):
273         """
274         Return bet and ids pairs.
275         """
276         if self._embedding_and_ids_eagerly:
277             return self._embedding_and_ids_eagerly
278         return self._embedding_and_ids_graph
279
```
--- 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
@@ -64,16 +64,21 @@
         # `None` dimension. This is because they have different shapes in
         # 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(
-                    initial_value=tf.zeros((1, self.output_dim)),
+                    # 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=tf.zeros((1, 1), dtype=tf.int64),
+                    initial_value=lambda: tf.zeros((1, 1), dtype=tf.int64),
                     shape=tf.TensorShape(None),
                     dtype=tf.int64,
                     trainable=False,
@@ -186,6 +191,13 @@
         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)

{"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@@ -64,16 +64,21 @@\n         # `None` dimension. This is because they have different shapes in\n         # different iterations.\n         # `tf.Variable` requires initial value if shape has `None` dimension.\n+        self._embedding_and_ids_graph = []\n+\n+    def _init_for_graph_mode(self):\n         self._embedding_and_ids_graph = [\n             EmbeddingAndIds(\n                 batch_embedding=tf.Variable(\n-                    initial_value=tf.zeros((1, self.output_dim)),\n+                    # In some cases, `tf.Variable` requires that initial value\n+                    # is callable.\n+                    initial_value=lambda: tf.zeros((1, self.output_dim)),\n                     shape=tf.TensorShape((None, self.output_dim)),\n                     dtype=tf.float32,\n                     trainable=True,\n                 ),\n                 batch_ids=tf.Variable(\n-                    initial_value=tf.zeros((1, 1), dtype=tf.int64),\n+                    initial_value=lambda: tf.zeros((1, 1), dtype=tf.int64),\n                     shape=tf.TensorShape(None),\n                     dtype=tf.int64,\n                     trainable=False,\n@@ -186,6 +191,13 @@\n         return batch_embedding\n \n     def call(self, input):\n+        if (\n+            self.tape\n+            and not tf.executing_eagerly()\n+            and not self._embedding_and_ids_graph\n+        ):\n+            self._init_for_graph_mode()\n+\n         input = tf.cast(input, tf.int64)\n         if isinstance(input, tf.SparseTensor):\n             return self._sparse_input_call(input)\n", "issue": "Do not save ElasticDL embedding's variables in the master\nIn current code, ElasticDL embedding layer create trainable variables and use them only for training under graph mode (i.e. under tf.function decorator). These variables are used for training and thus should only be created in workers. \r\n\r\nHowever, now ElasticDL embedding layer creates these variables in `__init__` function. The master (or PS) creates these variables when it creates a model instance, and then save these variables in `self._model`. \r\n\r\nI propose that ElasticDL embedding layer should not create these variables in `__init__` function. Instead, creating them before the training process only in worker is enough.\n", "before_files": [{"content": "import collections\n\nimport numpy as np\nimport tensorflow as tf\nfrom tensorflow.python.keras.utils import tf_utils\n\nfrom elasticdl.python.master.embedding_service import EmbeddingService\n\nEmbeddingAndIds = collections.namedtuple(\n    \"EmbeddingAndIds\", [\"batch_embedding\", \"batch_ids\"]\n)\n\n\nclass Embedding(tf.keras.layers.Layer):\n    \"\"\"\n    Input: indexes for the embedding entries with a shape of\n      (batch_size, input_length). Input can be either dense tensor\n      or SparseTensor.\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        If input is SparseTensor, mask_zero must be False.\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        If input is SparseTensor, combiner must set as a reduction op.\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        embedding_service_endpoint=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.supports_masking = mask_zero\n        self.input_length = input_length\n        self.combiner = combiner\n        self.embedding_service_endpoint = embedding_service_endpoint\n        self.tape = None\n        self.lookup_func = None\n\n        self._embedding_and_ids_eagerly = []\n\n        # BET's shape and ids' shape in `self._embedding_and_ids_graph` have\n        # `None` dimension. This is because they have different shapes in\n        # different iterations.\n        # `tf.Variable` requires initial value if shape has `None` dimension.\n        self._embedding_and_ids_graph = [\n            EmbeddingAndIds(\n                batch_embedding=tf.Variable(\n                    initial_value=tf.zeros((1, self.output_dim)),\n                    shape=tf.TensorShape((None, self.output_dim)),\n                    dtype=tf.float32,\n                    trainable=True,\n                ),\n                batch_ids=tf.Variable(\n                    initial_value=tf.zeros((1, 1), dtype=tf.int64),\n                    shape=tf.TensorShape(None),\n                    dtype=tf.int64,\n                    trainable=False,\n                ),\n            )\n        ]\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        ids = unique_ids.numpy()\n        keys = [Embedding.get_key([self._name, id]) for id in ids]\n        (\n            embedding_vectors,\n            unknown_keys_index,\n        ) = EmbeddingService.lookup_embedding(\n            keys=keys,\n            embedding_service_endpoint=self.embedding_service_endpoint,\n        )\n\n        if unknown_keys_index:\n            # Initialize unknown_keys' embedding vectors and write into Redis.\n            unknown_keys = [keys[index] for index in unknown_keys_index]\n            initializer = tf.keras.initializers.get(self.embedding_initializer)\n            embedding_vector_init = [\n                initializer(shape=[1, self.output_dim]).numpy()\n                for _ in unknown_keys\n            ]\n            embedding_vector_init = np.concatenate(\n                embedding_vector_init, axis=0\n            )\n            EmbeddingService.update_embedding(\n                keys=unknown_keys,\n                embedding_vectors=embedding_vector_init,\n                embedding_service_endpoint=self.embedding_service_endpoint,\n                set_if_not_exist=True,\n            )\n            # Lookup unknown_keys' embedding vectors\n            (\n                embedding_vectors_new,\n                unknown_keys_idx_new,\n            ) = EmbeddingService.lookup_embedding(\n                keys=unknown_keys,\n                embedding_service_endpoint=self.embedding_service_endpoint,\n            )\n            if unknown_keys_idx_new:\n                raise Exception(\n                    \"Update embedding vector: %s failed.\"\n                    % str(\n                        [unknown_keys[index] for index in unknown_keys_idx_new]\n                    )\n                )\n            for key_index, vector in zip(\n                unknown_keys_index, embedding_vectors_new\n            ):\n                embedding_vectors[key_index] = vector\n        embedding_vectors = np.concatenate(embedding_vectors, axis=0)\n        return embedding_vectors.reshape((len(keys), self.output_dim))\n\n    def _record_gradients(self, batch_embedding, ids):\n        if tf.executing_eagerly():\n            self.tape.watch(batch_embedding)\n            self._embedding_and_ids_eagerly.append(\n                EmbeddingAndIds(batch_embedding, ids)\n            )\n        else:\n            # In graph mode, assigning tensors to trainable variables is\n            # allowed and tape can record the gradients of trainable\n            # variables automatically.\n            embedding_and_ids = self._embedding_and_ids_graph[0]\n            embedding_and_ids.batch_embedding.assign(batch_embedding)\n            embedding_and_ids.batch_ids.assign(ids)\n            batch_embedding = embedding_and_ids.batch_embedding\n        return batch_embedding\n\n    def call(self, input):\n        input = tf.cast(input, tf.int64)\n        if isinstance(input, tf.SparseTensor):\n            return self._sparse_input_call(input)\n\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        # Gradient for `batch_embedding` is SparseTensor here due to\n        # `tf.gather` op. `tf.gather` accesses tensor slices, resulting in\n        # sparse tensor gradient.\n        batch_embedding = tf.py_function(\n            self.lookup_embedding, inp=[unique_ids], Tout=tf.float32\n        )\n        # TODO: use tf.cond rather than python if statement\n        if self.tape:\n            batch_embedding = self._record_gradients(batch_embedding, flat_ids)\n\n        outputs = tf.gather(batch_embedding, idx)\n        # tf.reshape does not support shape with None. Replace None with -1.\n        if ids.get_shape().rank == 2:\n            input_length = ids.get_shape()[1]\n            if input_length is None:\n                outputs.set_shape(shape=(None, None, self.output_dim))\n                return outputs\n            output_shape = (-1, input_length, self.output_dim)\n        else:\n            output_shape = ids.get_shape().concatenate(self.output_dim)\n        outputs = tf.reshape(outputs, output_shape)\n        # TODO: support combiner for dense input\n        return outputs\n\n    def _sparse_input_call(self, sparse_input):\n        if self.combiner not in [\"sum\", \"mean\", \"sqrtn\"]:\n            raise ValueError(\n                \"combiner must set sum, mean or sqrtn for sparse input\"\n            )\n        unique_ids, idx = tf.unique(sparse_input.values)\n        # Gradient for `batch_embedding` is dense tensor.\n        batch_embedding = tf.py_function(\n            self.lookup_embedding, inp=[unique_ids], Tout=tf.float32\n        )\n        # TODO: use tf.cond rather than python if statement\n        if self.tape:\n            batch_embedding = self._record_gradients(\n                batch_embedding, unique_ids\n            )\n\n        segment_ids = sparse_input.indices[:, 0]\n        if segment_ids.dtype != tf.int32:\n            segment_ids = tf.cast(segment_ids, tf.int32)\n\n        if self.combiner == \"sum\":\n            batch_embedding = tf.sparse.segment_sum(\n                batch_embedding, idx, segment_ids\n            )\n        elif self.combiner == \"mean\":\n            batch_embedding = tf.sparse.segment_mean(\n                batch_embedding, idx, segment_ids\n            )\n        elif self.combiner == \"sqrtn\":\n            batch_embedding = tf.sparse.segment_sqrt_n(\n                batch_embedding, idx, segment_ids\n            )\n        return batch_embedding\n\n    def compute_mask(self, inputs, mask=None):\n        if isinstance(input, tf.SparseTensor):\n            raise ValueError(\"SparseTensor inputs do not support mask_zero\")\n        if not self.supports_masking:\n            return None\n        return tf.math.not_equal(inputs, 0)\n\n    def reset(self):\n        self._embedding_and_ids_eagerly = []\n        self.tape = None\n\n    def set_tape(self, tape):\n        self.tape = tape\n\n    def set_endpoint(self, endpoint):\n        self.embedding_service_endpoint = endpoint\n\n    @property\n    def embedding_and_ids(self):\n        \"\"\"\n        Return bet and ids pairs.\n        \"\"\"\n        if self._embedding_and_ids_eagerly:\n            return self._embedding_and_ids_eagerly\n        return self._embedding_and_ids_graph\n", "path": "elasticdl/python/elasticdl/layers/embedding.py"}], "after_files": [{"content": "import collections\n\nimport numpy as np\nimport tensorflow as tf\nfrom tensorflow.python.keras.utils import tf_utils\n\nfrom elasticdl.python.master.embedding_service import EmbeddingService\n\nEmbeddingAndIds = collections.namedtuple(\n    \"EmbeddingAndIds\", [\"batch_embedding\", \"batch_ids\"]\n)\n\n\nclass Embedding(tf.keras.layers.Layer):\n    \"\"\"\n    Input: indexes for the embedding entries with a shape of\n      (batch_size, input_length). Input can be either dense tensor\n      or SparseTensor.\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        If input is SparseTensor, mask_zero must be False.\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        If input is SparseTensor, combiner must set as a reduction op.\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        embedding_service_endpoint=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.supports_masking = mask_zero\n        self.input_length = input_length\n        self.combiner = combiner\n        self.embedding_service_endpoint = embedding_service_endpoint\n        self.tape = None\n        self.lookup_func = None\n\n        self._embedding_and_ids_eagerly = []\n\n        # BET's shape and ids' shape in `self._embedding_and_ids_graph` have\n        # `None` dimension. This is because they have different shapes in\n        # different iterations.\n        # `tf.Variable` requires initial value if shape has `None` dimension.\n        self._embedding_and_ids_graph = []\n\n    def _init_for_graph_mode(self):\n        self._embedding_and_ids_graph = [\n            EmbeddingAndIds(\n                batch_embedding=tf.Variable(\n                    # In some cases, `tf.Variable` requires that initial value\n                    # is callable.\n                    initial_value=lambda: tf.zeros((1, self.output_dim)),\n                    shape=tf.TensorShape((None, self.output_dim)),\n                    dtype=tf.float32,\n                    trainable=True,\n                ),\n                batch_ids=tf.Variable(\n                    initial_value=lambda: tf.zeros((1, 1), dtype=tf.int64),\n                    shape=tf.TensorShape(None),\n                    dtype=tf.int64,\n                    trainable=False,\n                ),\n            )\n        ]\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        ids = unique_ids.numpy()\n        keys = [Embedding.get_key([self._name, id]) for id in ids]\n        (\n            embedding_vectors,\n            unknown_keys_index,\n        ) = EmbeddingService.lookup_embedding(\n            keys=keys,\n            embedding_service_endpoint=self.embedding_service_endpoint,\n        )\n\n        if unknown_keys_index:\n            # Initialize unknown_keys' embedding vectors and write into Redis.\n            unknown_keys = [keys[index] for index in unknown_keys_index]\n            initializer = tf.keras.initializers.get(self.embedding_initializer)\n            embedding_vector_init = [\n                initializer(shape=[1, self.output_dim]).numpy()\n                for _ in unknown_keys\n            ]\n            embedding_vector_init = np.concatenate(\n                embedding_vector_init, axis=0\n            )\n            EmbeddingService.update_embedding(\n                keys=unknown_keys,\n                embedding_vectors=embedding_vector_init,\n                embedding_service_endpoint=self.embedding_service_endpoint,\n                set_if_not_exist=True,\n            )\n            # Lookup unknown_keys' embedding vectors\n            (\n                embedding_vectors_new,\n                unknown_keys_idx_new,\n            ) = EmbeddingService.lookup_embedding(\n                keys=unknown_keys,\n                embedding_service_endpoint=self.embedding_service_endpoint,\n            )\n            if unknown_keys_idx_new:\n                raise Exception(\n                    \"Update embedding vector: %s failed.\"\n                    % str(\n                        [unknown_keys[index] for index in unknown_keys_idx_new]\n                    )\n                )\n            for key_index, vector in zip(\n                unknown_keys_index, embedding_vectors_new\n            ):\n                embedding_vectors[key_index] = vector\n        embedding_vectors = np.concatenate(embedding_vectors, axis=0)\n        return embedding_vectors.reshape((len(keys), self.output_dim))\n\n    def _record_gradients(self, batch_embedding, ids):\n        if tf.executing_eagerly():\n            self.tape.watch(batch_embedding)\n            self._embedding_and_ids_eagerly.append(\n                EmbeddingAndIds(batch_embedding, ids)\n            )\n        else:\n            # In graph mode, assigning tensors to trainable variables is\n            # allowed and tape can record the gradients of trainable\n            # variables automatically.\n            embedding_and_ids = self._embedding_and_ids_graph[0]\n            embedding_and_ids.batch_embedding.assign(batch_embedding)\n            embedding_and_ids.batch_ids.assign(ids)\n            batch_embedding = embedding_and_ids.batch_embedding\n        return batch_embedding\n\n    def call(self, input):\n        if (\n            self.tape\n            and not tf.executing_eagerly()\n            and not self._embedding_and_ids_graph\n        ):\n            self._init_for_graph_mode()\n\n        input = tf.cast(input, tf.int64)\n        if isinstance(input, tf.SparseTensor):\n            return self._sparse_input_call(input)\n\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        # Gradient for `batch_embedding` is SparseTensor here due to\n        # `tf.gather` op. `tf.gather` accesses tensor slices, resulting in\n        # sparse tensor gradient.\n        batch_embedding = tf.py_function(\n            self.lookup_embedding, inp=[unique_ids], Tout=tf.float32\n        )\n        # TODO: use tf.cond rather than python if statement\n        if self.tape:\n            batch_embedding = self._record_gradients(batch_embedding, flat_ids)\n\n        outputs = tf.gather(batch_embedding, idx)\n        # tf.reshape does not support shape with None. Replace None with -1.\n        if ids.get_shape().rank == 2:\n            input_length = ids.get_shape()[1]\n            if input_length is None:\n                outputs.set_shape(shape=(None, None, self.output_dim))\n                return outputs\n            output_shape = (-1, input_length, self.output_dim)\n        else:\n            output_shape = ids.get_shape().concatenate(self.output_dim)\n        outputs = tf.reshape(outputs, output_shape)\n        # TODO: support combiner for dense input\n        return outputs\n\n    def _sparse_input_call(self, sparse_input):\n        if self.combiner not in [\"sum\", \"mean\", \"sqrtn\"]:\n            raise ValueError(\n                \"combiner must set sum, mean or sqrtn for sparse input\"\n            )\n        unique_ids, idx = tf.unique(sparse_input.values)\n        # Gradient for `batch_embedding` is dense tensor.\n        batch_embedding = tf.py_function(\n            self.lookup_embedding, inp=[unique_ids], Tout=tf.float32\n        )\n        # TODO: use tf.cond rather than python if statement\n        if self.tape:\n            batch_embedding = self._record_gradients(\n                batch_embedding, unique_ids\n            )\n\n        segment_ids = sparse_input.indices[:, 0]\n        if segment_ids.dtype != tf.int32:\n            segment_ids = tf.cast(segment_ids, tf.int32)\n\n        if self.combiner == \"sum\":\n            batch_embedding = tf.sparse.segment_sum(\n                batch_embedding, idx, segment_ids\n            )\n        elif self.combiner == \"mean\":\n            batch_embedding = tf.sparse.segment_mean(\n                batch_embedding, idx, segment_ids\n            )\n        elif self.combiner == \"sqrtn\":\n            batch_embedding = tf.sparse.segment_sqrt_n(\n                batch_embedding, idx, segment_ids\n            )\n        return batch_embedding\n\n    def compute_mask(self, inputs, mask=None):\n        if isinstance(input, tf.SparseTensor):\n            raise ValueError(\"SparseTensor inputs do not support mask_zero\")\n        if not self.supports_masking:\n            return None\n        return tf.math.not_equal(inputs, 0)\n\n    def reset(self):\n        self._embedding_and_ids_eagerly = []\n        self.tape = None\n\n    def set_tape(self, tape):\n        self.tape = tape\n\n    def set_endpoint(self, endpoint):\n        self.embedding_service_endpoint = endpoint\n\n    @property\n    def embedding_and_ids(self):\n        \"\"\"\n        Return bet and ids pairs.\n        \"\"\"\n        if self._embedding_and_ids_eagerly:\n            return self._embedding_and_ids_eagerly\n        return self._embedding_and_ids_graph\n", "path": "elasticdl/python/elasticdl/layers/embedding.py"}]}