We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Embedding Bag Fails GPU Testing
It looks like our new layer is failing for GPU tests. 

https://github.com/tensorflow/addons/pull/2501

https://source.cloud.google.com/results/invocations/42c93476-b0f7-4192-9da5-0dac23c4c92b/targets/tensorflow_addons%2Fubuntu%2Fgpu%2Fpy3%2Fpresubmit/log

@Rocketknight1 when time allows could you please take a look? We can trigger GPU testing on any PR you submit.


e.g.
```
tensorflow_addons/layers/tests/embedding_bag_test.py:120:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

a = array([[0.40625, 0.40625, 0.40625, ..., 0.40625, 0.40625, 0.40625],
       [0.25   , 0.25   , 0.25   , ..., 0.25   , 0...0.34375, 0.34375, ..., 0.34375, 0.34375, 0.34375],
       [0.21875, 0.21875, 0.21875, ..., 0.21875, 0.21875, 0.21875]])
b = array([[13., 13., 13., ..., 13., 13., 13.],
       [ 8.,  8.,  8., ...,  8.,  8.,  8.],
       [ 7.,  7.,  7., ...,  7... 7.,  7., ...,  7.,  7.,  7.],
       [11., 11., 11., ..., 11., 11., 11.],
       [ 7.,  7.,  7., ...,  7.,  7.,  7.]])
rtol = 1e-06, atol = 1e-06, float_rtol = 1e-06, float_atol = 1e-06
half_rtol = 0.001, half_atol = 0.001, bfloat16_rtol = 0.01, bfloat16_atol = 0.01

    def assert_allclose_according_to_type(
        a,
        b,
        rtol=1e-6,
        atol=1e-6,
        float_rtol=1e-6,
        float_atol=1e-6,
        half_rtol=1e-3,
        half_atol=1e-3,
        bfloat16_rtol=1e-2,
        bfloat16_atol=1e-2,
    ):
        """
        Similar to tf.test.TestCase.assertAllCloseAccordingToType()
        but this doesn't need a subclassing to run.
        """
        a = np.array(a)
        b = np.array(b)
        # types with lower tol are put later to overwrite previous ones.
        if (
            a.dtype == np.float32
            or b.dtype == np.float32
            or a.dtype == np.complex64
            or b.dtype == np.complex64
        ):
            rtol = max(rtol, float_rtol)
            atol = max(atol, float_atol)
        if a.dtype == np.float16 or b.dtype == np.float16:
            rtol = max(rtol, half_rtol)
            atol = max(atol, half_atol)
        if a.dtype == tf.bfloat16.as_numpy_dtype or b.dtype == tf.bfloat16.as_numpy_dtype:
            rtol = max(rtol, bfloat16_rtol)
            atol = max(atol, bfloat16_atol)

>       np.testing.assert_allclose(a, b, rtol=rtol, atol=atol)
E       AssertionError: 
E       Not equal to tolerance rtol=1e-06, atol=1e-06
E       
E       Mismatched elements: 1008 / 1008 (100%)
E       Max absolute difference: 14.53125
E       Max relative difference: 0.96875
E        x: array([[0.40625, 0.40625, 0.40625, ..., 0.40625, 0.40625, 0.40625],
E              [0.25   , 0.25   , 0.25   , ..., 0.25   , 0.25   , 0.25   ],
E              [0.21875, 0.21875, 0.21875, ..., 0.21875, 0.21875, 0.21875],...
E        y: array([[13., 13., 13., ..., 13., 13., 13.],
E              [ 8.,  8.,  8., ...,  8.,  8.,  8.],
E              [ 7.,  7.,  7., ...,  7.,  7.,  7.],...
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `tensorflow_addons/layers/embedding_bag.py`
Content:
```
1 # Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # ==============================================================================
15 
16 import tensorflow as tf
17 from typeguard import typechecked
18 
19 from tensorflow_addons.utils.types import Constraint, Initializer, Regularizer
20 from tensorflow_addons.utils.resource_loader import LazySO
21 
22 _embedding_bag_so = LazySO("custom_ops/layers/_embedding_bag_ops.so")
23 
24 
25 def _embedding_bag(
26     indices,
27     params,
28     weights=None,
29     combiner="sum",
30     name=None,
31 ):
32     """EmbeddingBag computation.
33 
34     See [PyTorch op](https://pytorch.org/docs/stable/generated/torch.nn.EmbeddingBag.html).
35 
36     Equivalent to tf.gather() followed by tf.reduce_{sum,mean}() across the last dimension, with optional
37     weights. Fusing these into a single op has massive benefits for execution speed and particularly
38     memory usage, as the intermediate output of the gather never needs to be materialized.
39 
40     Args:
41       indices: An int32 or int64 `Tensor` of the indices to gather from
42           `params`. Must be at least 2-dimensional, as the last dimension
43           will be summed out. Maximum value must be less than params.shape[0].
44       params: A float32 `Tensor` from which to gather params. Must be rank 2.
45       weights: A float32 `Tensor` of weights which will be applied to each of
46           the gathered embedding vectors before the sum step.
47       name: A name for the operation (optional).
48 
49     Returns:
50       A `Tensor` of the format specified by `data_format`.
51     """
52     if weights is None:
53         weights = tf.ones_like(indices, dtype=params.dtype)
54     elif combiner != "sum":
55         raise RuntimeError(
56             "Combiner mode must be 'sum' when weights are supplied to EmbeddingBag!"
57         )
58 
59     return _embedding_bag_so.ops.addons_embedding_bag(
60         indices, params, weights, combiner=combiner.upper(), name=name
61     )
62 
63 
64 @tf.RegisterGradient("Addons>EmbeddingBag")
65 def _embedding_bag_grad(op, grads):
66     indices, params, weights = op.inputs[:3]
67     combiner = op.get_attr("combiner")
68     value_grads, weight_grads = _embedding_bag_so.ops.addons_embedding_bag_grad(
69         indices, params, weights, grads, combiner=combiner
70     )
71     return [None, value_grads, weight_grads]
72 
73 
74 @tf.keras.utils.register_keras_serializable(package="Addons")
75 class EmbeddingBag(tf.keras.layers.Layer):
76     """EmbeddingBag Layer.
77 
78     See [PyTorch op](https://pytorch.org/docs/stable/generated/torch.nn.EmbeddingBag.html).
79 
80     Equivalent to tf.gather() followed by tf.reduce_sum() across the last dimension, with optional
81     weights. Fusing these into a single op has massive benefits for execution speed and particularly
82     memory usage, as the intermediate output of the gather never needs to be materialized.
83 
84     Input Shapes:
85       indices: An int32 or int64 `Tensor` of the indices to gather from
86           `params`. Must be at least 2-dimensional, as the last dimension
87           will be summed out. Maximum value must be less than params.shape[0].
88       params: A float32 `Tensor` from which to gather params. Must be rank 2.
89       weights: A float32 `Tensor` of weights which will be applied to each of
90           the gathered embedding vectors before the sum step.
91 
92     Output shape:
93         indices.shape[:-1], params.shape[-1]
94     """
95 
96     @typechecked
97     def __init__(
98         self,
99         input_dim: int,
100         output_dim: int,
101         embeddings_initializer: Initializer = "uniform",
102         embeddings_regularizer: Regularizer = None,
103         embeddings_constraint: Constraint = None,
104         mask_zero: bool = False,
105         combiner: str = "sum",
106         **kwargs,
107     ):
108         super(EmbeddingBag, self).__init__(**kwargs)
109         if input_dim <= 0 or output_dim <= 0:
110             raise ValueError(
111                 "Both `input_dim` and `output_dim` should be positive, "
112                 "found input_dim {} and output_dim {}".format(input_dim, output_dim)
113             )
114         self.input_dim = input_dim
115         self.output_dim = output_dim
116         self.embeddings_initializer = tf.keras.initializers.get(embeddings_initializer)
117         self.embeddings_regularizer = tf.keras.regularizers.get(embeddings_regularizer)
118         self.embeddings_constraint = tf.keras.constraints.get(embeddings_constraint)
119         self.mask_zero = mask_zero
120         self.supports_masking = mask_zero
121         self.combiner = combiner
122 
123     def build(self, input_shape):
124         self.embeddings = self.add_weight(
125             shape=(self.input_dim, self.output_dim),
126             name="embeddings",
127             initializer=self.embeddings_initializer,
128             regularizer=self.embeddings_regularizer,
129             constraint=self.embeddings_constraint,
130         )
131         self.built = True
132 
133     def call(self, indices, weights=None):
134         return _embedding_bag(indices, self.embeddings, weights, combiner=self.combiner)
135 
136     def get_config(self):
137         config = {
138             "input_dim": self.input_dim,
139             "output_dim": self.output_dim,
140             "embeddings_initializer": tf.keras.initializers.serialize(
141                 self.embeddings_initializer
142             ),
143             "embeddings_regularizer": tf.keras.regularizers.serialize(
144                 self.embeddings_regularizer
145             ),
146             "embeddings_constraint": tf.keras.constraints.serialize(
147                 self.embeddings_constraint
148             ),
149             "mask_zero": self.mask_zero,
150             "input_length": self.input_length,
151             "combiner": self.combiner,
152         }
153         base_config = super(EmbeddingBag, self).get_config()
154         return dict(list(base_config.items()) + list(config.items()))
155
```
--- 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/tensorflow_addons/layers/embedding_bag.py b/tensorflow_addons/layers/embedding_bag.py
--- a/tensorflow_addons/layers/embedding_bag.py
+++ b/tensorflow_addons/layers/embedding_bag.py
@@ -49,8 +49,13 @@
     Returns:
       A `Tensor` of the format specified by `data_format`.
     """
-    if weights is None:
+    if weights is None and combiner == "sum":
         weights = tf.ones_like(indices, dtype=params.dtype)
+    elif weights is None and combiner == "mean":
+        weights = tf.ones_like(indices, dtype=params.dtype) / tf.cast(
+            tf.shape(indices)[1], params.dtype
+        )
+        combiner = "sum"
     elif combiner != "sum":
         raise RuntimeError(
             "Combiner mode must be 'sum' when weights are supplied to EmbeddingBag!"

{"golden_diff": "diff --git a/tensorflow_addons/layers/embedding_bag.py b/tensorflow_addons/layers/embedding_bag.py\n--- a/tensorflow_addons/layers/embedding_bag.py\n+++ b/tensorflow_addons/layers/embedding_bag.py\n@@ -49,8 +49,13 @@\n     Returns:\n       A `Tensor` of the format specified by `data_format`.\n     \"\"\"\n-    if weights is None:\n+    if weights is None and combiner == \"sum\":\n         weights = tf.ones_like(indices, dtype=params.dtype)\n+    elif weights is None and combiner == \"mean\":\n+        weights = tf.ones_like(indices, dtype=params.dtype) / tf.cast(\n+            tf.shape(indices)[1], params.dtype\n+        )\n+        combiner = \"sum\"\n     elif combiner != \"sum\":\n         raise RuntimeError(\n             \"Combiner mode must be 'sum' when weights are supplied to EmbeddingBag!\"\n", "issue": "Embedding Bag Fails GPU Testing\nIt looks like our new layer is failing for GPU tests. \r\n\r\nhttps://github.com/tensorflow/addons/pull/2501\r\n\r\nhttps://source.cloud.google.com/results/invocations/42c93476-b0f7-4192-9da5-0dac23c4c92b/targets/tensorflow_addons%2Fubuntu%2Fgpu%2Fpy3%2Fpresubmit/log\r\n\r\n@Rocketknight1 when time allows could you please take a look? We can trigger GPU testing on any PR you submit.\r\n\r\n\r\ne.g.\r\n```\r\ntensorflow_addons/layers/tests/embedding_bag_test.py:120:\r\n_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _\r\n\r\na = array([[0.40625, 0.40625, 0.40625, ..., 0.40625, 0.40625, 0.40625],\r\n       [0.25   , 0.25   , 0.25   , ..., 0.25   , 0...0.34375, 0.34375, ..., 0.34375, 0.34375, 0.34375],\r\n       [0.21875, 0.21875, 0.21875, ..., 0.21875, 0.21875, 0.21875]])\r\nb = array([[13., 13., 13., ..., 13., 13., 13.],\r\n       [ 8.,  8.,  8., ...,  8.,  8.,  8.],\r\n       [ 7.,  7.,  7., ...,  7... 7.,  7., ...,  7.,  7.,  7.],\r\n       [11., 11., 11., ..., 11., 11., 11.],\r\n       [ 7.,  7.,  7., ...,  7.,  7.,  7.]])\r\nrtol = 1e-06, atol = 1e-06, float_rtol = 1e-06, float_atol = 1e-06\r\nhalf_rtol = 0.001, half_atol = 0.001, bfloat16_rtol = 0.01, bfloat16_atol = 0.01\r\n\r\n    def assert_allclose_according_to_type(\r\n        a,\r\n        b,\r\n        rtol=1e-6,\r\n        atol=1e-6,\r\n        float_rtol=1e-6,\r\n        float_atol=1e-6,\r\n        half_rtol=1e-3,\r\n        half_atol=1e-3,\r\n        bfloat16_rtol=1e-2,\r\n        bfloat16_atol=1e-2,\r\n    ):\r\n        \"\"\"\r\n        Similar to tf.test.TestCase.assertAllCloseAccordingToType()\r\n        but this doesn't need a subclassing to run.\r\n        \"\"\"\r\n        a = np.array(a)\r\n        b = np.array(b)\r\n        # types with lower tol are put later to overwrite previous ones.\r\n        if (\r\n            a.dtype == np.float32\r\n            or b.dtype == np.float32\r\n            or a.dtype == np.complex64\r\n            or b.dtype == np.complex64\r\n        ):\r\n            rtol = max(rtol, float_rtol)\r\n            atol = max(atol, float_atol)\r\n        if a.dtype == np.float16 or b.dtype == np.float16:\r\n            rtol = max(rtol, half_rtol)\r\n            atol = max(atol, half_atol)\r\n        if a.dtype == tf.bfloat16.as_numpy_dtype or b.dtype == tf.bfloat16.as_numpy_dtype:\r\n            rtol = max(rtol, bfloat16_rtol)\r\n            atol = max(atol, bfloat16_atol)\r\n\r\n>       np.testing.assert_allclose(a, b, rtol=rtol, atol=atol)\r\nE       AssertionError: \r\nE       Not equal to tolerance rtol=1e-06, atol=1e-06\r\nE       \r\nE       Mismatched elements: 1008 / 1008 (100%)\r\nE       Max absolute difference: 14.53125\r\nE       Max relative difference: 0.96875\r\nE        x: array([[0.40625, 0.40625, 0.40625, ..., 0.40625, 0.40625, 0.40625],\r\nE              [0.25   , 0.25   , 0.25   , ..., 0.25   , 0.25   , 0.25   ],\r\nE              [0.21875, 0.21875, 0.21875, ..., 0.21875, 0.21875, 0.21875],...\r\nE        y: array([[13., 13., 13., ..., 13., 13., 13.],\r\nE              [ 8.,  8.,  8., ...,  8.,  8.,  8.],\r\nE              [ 7.,  7.,  7., ...,  7.,  7.,  7.],...\r\n```\n", "before_files": [{"content": "# Copyright 2021 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\nimport tensorflow as tf\nfrom typeguard import typechecked\n\nfrom tensorflow_addons.utils.types import Constraint, Initializer, Regularizer\nfrom tensorflow_addons.utils.resource_loader import LazySO\n\n_embedding_bag_so = LazySO(\"custom_ops/layers/_embedding_bag_ops.so\")\n\n\ndef _embedding_bag(\n    indices,\n    params,\n    weights=None,\n    combiner=\"sum\",\n    name=None,\n):\n    \"\"\"EmbeddingBag computation.\n\n    See [PyTorch op](https://pytorch.org/docs/stable/generated/torch.nn.EmbeddingBag.html).\n\n    Equivalent to tf.gather() followed by tf.reduce_{sum,mean}() across the last dimension, with optional\n    weights. Fusing these into a single op has massive benefits for execution speed and particularly\n    memory usage, as the intermediate output of the gather never needs to be materialized.\n\n    Args:\n      indices: An int32 or int64 `Tensor` of the indices to gather from\n          `params`. Must be at least 2-dimensional, as the last dimension\n          will be summed out. Maximum value must be less than params.shape[0].\n      params: A float32 `Tensor` from which to gather params. Must be rank 2.\n      weights: A float32 `Tensor` of weights which will be applied to each of\n          the gathered embedding vectors before the sum step.\n      name: A name for the operation (optional).\n\n    Returns:\n      A `Tensor` of the format specified by `data_format`.\n    \"\"\"\n    if weights is None:\n        weights = tf.ones_like(indices, dtype=params.dtype)\n    elif combiner != \"sum\":\n        raise RuntimeError(\n            \"Combiner mode must be 'sum' when weights are supplied to EmbeddingBag!\"\n        )\n\n    return _embedding_bag_so.ops.addons_embedding_bag(\n        indices, params, weights, combiner=combiner.upper(), name=name\n    )\n\n\n@tf.RegisterGradient(\"Addons>EmbeddingBag\")\ndef _embedding_bag_grad(op, grads):\n    indices, params, weights = op.inputs[:3]\n    combiner = op.get_attr(\"combiner\")\n    value_grads, weight_grads = _embedding_bag_so.ops.addons_embedding_bag_grad(\n        indices, params, weights, grads, combiner=combiner\n    )\n    return [None, value_grads, weight_grads]\n\n\n@tf.keras.utils.register_keras_serializable(package=\"Addons\")\nclass EmbeddingBag(tf.keras.layers.Layer):\n    \"\"\"EmbeddingBag Layer.\n\n    See [PyTorch op](https://pytorch.org/docs/stable/generated/torch.nn.EmbeddingBag.html).\n\n    Equivalent to tf.gather() followed by tf.reduce_sum() across the last dimension, with optional\n    weights. Fusing these into a single op has massive benefits for execution speed and particularly\n    memory usage, as the intermediate output of the gather never needs to be materialized.\n\n    Input Shapes:\n      indices: An int32 or int64 `Tensor` of the indices to gather from\n          `params`. Must be at least 2-dimensional, as the last dimension\n          will be summed out. Maximum value must be less than params.shape[0].\n      params: A float32 `Tensor` from which to gather params. Must be rank 2.\n      weights: A float32 `Tensor` of weights which will be applied to each of\n          the gathered embedding vectors before the sum step.\n\n    Output shape:\n        indices.shape[:-1], params.shape[-1]\n    \"\"\"\n\n    @typechecked\n    def __init__(\n        self,\n        input_dim: int,\n        output_dim: int,\n        embeddings_initializer: Initializer = \"uniform\",\n        embeddings_regularizer: Regularizer = None,\n        embeddings_constraint: Constraint = None,\n        mask_zero: bool = False,\n        combiner: str = \"sum\",\n        **kwargs,\n    ):\n        super(EmbeddingBag, self).__init__(**kwargs)\n        if input_dim <= 0 or output_dim <= 0:\n            raise ValueError(\n                \"Both `input_dim` and `output_dim` should be positive, \"\n                \"found input_dim {} and output_dim {}\".format(input_dim, output_dim)\n            )\n        self.input_dim = input_dim\n        self.output_dim = output_dim\n        self.embeddings_initializer = tf.keras.initializers.get(embeddings_initializer)\n        self.embeddings_regularizer = tf.keras.regularizers.get(embeddings_regularizer)\n        self.embeddings_constraint = tf.keras.constraints.get(embeddings_constraint)\n        self.mask_zero = mask_zero\n        self.supports_masking = mask_zero\n        self.combiner = combiner\n\n    def build(self, input_shape):\n        self.embeddings = self.add_weight(\n            shape=(self.input_dim, self.output_dim),\n            name=\"embeddings\",\n            initializer=self.embeddings_initializer,\n            regularizer=self.embeddings_regularizer,\n            constraint=self.embeddings_constraint,\n        )\n        self.built = True\n\n    def call(self, indices, weights=None):\n        return _embedding_bag(indices, self.embeddings, weights, combiner=self.combiner)\n\n    def get_config(self):\n        config = {\n            \"input_dim\": self.input_dim,\n            \"output_dim\": self.output_dim,\n            \"embeddings_initializer\": tf.keras.initializers.serialize(\n                self.embeddings_initializer\n            ),\n            \"embeddings_regularizer\": tf.keras.regularizers.serialize(\n                self.embeddings_regularizer\n            ),\n            \"embeddings_constraint\": tf.keras.constraints.serialize(\n                self.embeddings_constraint\n            ),\n            \"mask_zero\": self.mask_zero,\n            \"input_length\": self.input_length,\n            \"combiner\": self.combiner,\n        }\n        base_config = super(EmbeddingBag, self).get_config()\n        return dict(list(base_config.items()) + list(config.items()))\n", "path": "tensorflow_addons/layers/embedding_bag.py"}], "after_files": [{"content": "# Copyright 2021 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\nimport tensorflow as tf\nfrom typeguard import typechecked\n\nfrom tensorflow_addons.utils.types import Constraint, Initializer, Regularizer\nfrom tensorflow_addons.utils.resource_loader import LazySO\n\n_embedding_bag_so = LazySO(\"custom_ops/layers/_embedding_bag_ops.so\")\n\n\ndef _embedding_bag(\n    indices,\n    params,\n    weights=None,\n    combiner=\"sum\",\n    name=None,\n):\n    \"\"\"EmbeddingBag computation.\n\n    See [PyTorch op](https://pytorch.org/docs/stable/generated/torch.nn.EmbeddingBag.html).\n\n    Equivalent to tf.gather() followed by tf.reduce_{sum,mean}() across the last dimension, with optional\n    weights. Fusing these into a single op has massive benefits for execution speed and particularly\n    memory usage, as the intermediate output of the gather never needs to be materialized.\n\n    Args:\n      indices: An int32 or int64 `Tensor` of the indices to gather from\n          `params`. Must be at least 2-dimensional, as the last dimension\n          will be summed out. Maximum value must be less than params.shape[0].\n      params: A float32 `Tensor` from which to gather params. Must be rank 2.\n      weights: A float32 `Tensor` of weights which will be applied to each of\n          the gathered embedding vectors before the sum step.\n      name: A name for the operation (optional).\n\n    Returns:\n      A `Tensor` of the format specified by `data_format`.\n    \"\"\"\n    if weights is None and combiner == \"sum\":\n        weights = tf.ones_like(indices, dtype=params.dtype)\n    elif weights is None and combiner == \"mean\":\n        weights = tf.ones_like(indices, dtype=params.dtype) / tf.cast(\n            tf.shape(indices)[1], params.dtype\n        )\n        combiner = \"sum\"\n    elif combiner != \"sum\":\n        raise RuntimeError(\n            \"Combiner mode must be 'sum' when weights are supplied to EmbeddingBag!\"\n        )\n\n    return _embedding_bag_so.ops.addons_embedding_bag(\n        indices, params, weights, combiner=combiner.upper(), name=name\n    )\n\n\n@tf.RegisterGradient(\"Addons>EmbeddingBag\")\ndef _embedding_bag_grad(op, grads):\n    indices, params, weights = op.inputs[:3]\n    combiner = op.get_attr(\"combiner\")\n    value_grads, weight_grads = _embedding_bag_so.ops.addons_embedding_bag_grad(\n        indices, params, weights, grads, combiner=combiner\n    )\n    return [None, value_grads, weight_grads]\n\n\n@tf.keras.utils.register_keras_serializable(package=\"Addons\")\nclass EmbeddingBag(tf.keras.layers.Layer):\n    \"\"\"EmbeddingBag Layer.\n\n    See [PyTorch op](https://pytorch.org/docs/stable/generated/torch.nn.EmbeddingBag.html).\n\n    Equivalent to tf.gather() followed by tf.reduce_sum() across the last dimension, with optional\n    weights. Fusing these into a single op has massive benefits for execution speed and particularly\n    memory usage, as the intermediate output of the gather never needs to be materialized.\n\n    Input Shapes:\n      indices: An int32 or int64 `Tensor` of the indices to gather from\n          `params`. Must be at least 2-dimensional, as the last dimension\n          will be summed out. Maximum value must be less than params.shape[0].\n      params: A float32 `Tensor` from which to gather params. Must be rank 2.\n      weights: A float32 `Tensor` of weights which will be applied to each of\n          the gathered embedding vectors before the sum step.\n\n    Output shape:\n        indices.shape[:-1], params.shape[-1]\n    \"\"\"\n\n    @typechecked\n    def __init__(\n        self,\n        input_dim: int,\n        output_dim: int,\n        embeddings_initializer: Initializer = \"uniform\",\n        embeddings_regularizer: Regularizer = None,\n        embeddings_constraint: Constraint = None,\n        mask_zero: bool = False,\n        combiner: str = \"sum\",\n        **kwargs,\n    ):\n        super(EmbeddingBag, self).__init__(**kwargs)\n        if input_dim <= 0 or output_dim <= 0:\n            raise ValueError(\n                \"Both `input_dim` and `output_dim` should be positive, \"\n                \"found input_dim {} and output_dim {}\".format(input_dim, output_dim)\n            )\n        self.input_dim = input_dim\n        self.output_dim = output_dim\n        self.embeddings_initializer = tf.keras.initializers.get(embeddings_initializer)\n        self.embeddings_regularizer = tf.keras.regularizers.get(embeddings_regularizer)\n        self.embeddings_constraint = tf.keras.constraints.get(embeddings_constraint)\n        self.mask_zero = mask_zero\n        self.supports_masking = mask_zero\n        self.combiner = combiner\n\n    def build(self, input_shape):\n        self.embeddings = self.add_weight(\n            shape=(self.input_dim, self.output_dim),\n            name=\"embeddings\",\n            initializer=self.embeddings_initializer,\n            regularizer=self.embeddings_regularizer,\n            constraint=self.embeddings_constraint,\n        )\n        self.built = True\n\n    def call(self, indices, weights=None):\n        return _embedding_bag(indices, self.embeddings, weights, combiner=self.combiner)\n\n    def get_config(self):\n        config = {\n            \"input_dim\": self.input_dim,\n            \"output_dim\": self.output_dim,\n            \"embeddings_initializer\": tf.keras.initializers.serialize(\n                self.embeddings_initializer\n            ),\n            \"embeddings_regularizer\": tf.keras.regularizers.serialize(\n                self.embeddings_regularizer\n            ),\n            \"embeddings_constraint\": tf.keras.constraints.serialize(\n                self.embeddings_constraint\n            ),\n            \"mask_zero\": self.mask_zero,\n            \"input_length\": self.input_length,\n            \"combiner\": self.combiner,\n        }\n        base_config = super(EmbeddingBag, self).get_config()\n        return dict(list(base_config.items()) + list(config.items()))\n", "path": "tensorflow_addons/layers/embedding_bag.py"}]}