Repository: tensorflow/addons

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.],...
```


Discussion:
This is very odd - I included the tests with the PR and they passed at the time. I'll try to figure out what's going on when I get a chance this week.
This bug seems to only be triggered when the combiner mode is set to 'mean' - that narrows it down a lot, and means the CUDA is probably fine. I'm investigating further but hopefully it's just a small issue. I just don't know how it made it through the checks when I first merged it!
Thanks @Rocketknight1 for the investigation! The way it made it through checks is our fault. In order to run GPU tests we have to attach a label which triggers Google to test with a GPU instance. There's no GPU instances available for Github CI so this is much appreciated but it does cause some things to slip through the cracks occassionally.
We used the kokoro label on the PR but probably not on the latest commit.

Resolving patch:
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 @@ def _embedding_bag(
     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!"


Test patch:
diff --git a/tensorflow_addons/layers/tests/embedding_bag_test.py b/tensorflow_addons/layers/tests/embedding_bag_test.py
--- a/tensorflow_addons/layers/tests/embedding_bag_test.py
+++ b/tensorflow_addons/layers/tests/embedding_bag_test.py
@@ -60,7 +60,9 @@ def test_forward(input_shape, input_dim, dtype, indices_dtype, combiner):
         indices,
         weights,
     )
-    test_utils.assert_allclose_according_to_type(expected, output)
+    test_utils.assert_allclose_according_to_type(
+        expected, output, half_rtol=1e-2, half_atol=1e-2
+    )
 
 
 @pytest.mark.with_device(["cpu", "gpu"])
@@ -69,8 +71,8 @@ def test_forward(input_shape, input_dim, dtype, indices_dtype, combiner):
 @pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64])
 @pytest.mark.parametrize("indices_dtype", [np.int32, np.int64])
 @pytest.mark.parametrize("combiner", ["sum", "mean"])
-@pytest.mark.parametrize("graph_mode", [True, False])
-def test_backward(input_shape, input_dim, dtype, indices_dtype, combiner, graph_mode):
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_backward(input_shape, input_dim, dtype, indices_dtype, combiner):
     indices = np.random.randint(low=0, high=input_dim, size=input_shape).astype(
         indices_dtype
     )
@@ -85,10 +87,7 @@ def test_backward(input_shape, input_dim, dtype, indices_dtype, combiner, graph_
     if weights is not None:
         weights = tf.convert_to_tensor(weights)
 
-    if graph_mode:
-        embedding_bag_fn = tf.function(_embedding_bag)
-    else:
-        embedding_bag_fn = _embedding_bag
+    embedding_bag_fn = tf.function(_embedding_bag)
 
     if combiner == "sum":
         with tf.GradientTape(persistent=True) as tape:
@@ -102,10 +101,11 @@ def test_backward(input_shape, input_dim, dtype, indices_dtype, combiner, graph_
         test_utils.assert_allclose_according_to_type(
             tf.convert_to_tensor(expected_grads[0]),
             tf.convert_to_tensor(grads[0]),
+            half_rtol=1e-2,
+            half_atol=1e-2,
         )
         test_utils.assert_allclose_according_to_type(
-            expected_grads[1],
-            grads[1],
+            expected_grads[1], grads[1], half_rtol=1e-2, half_atol=1e-2
         )
     else:
         with tf.GradientTape(persistent=True) as tape:
@@ -117,5 +117,8 @@ def test_backward(input_shape, input_dim, dtype, indices_dtype, combiner, graph_
         expected_grads = tape.gradient(expected, [params])
         # Gather returns sparse IndexedSlices so we have to sum them together.
         test_utils.assert_allclose_according_to_type(
-            tf.convert_to_tensor(expected_grads[0]), tf.convert_to_tensor(grads[0])
+            tf.convert_to_tensor(expected_grads[0]),
+            tf.convert_to_tensor(grads[0]),
+            half_rtol=1e-2,
+            half_atol=1e-2,
         )