Repository: tensorflow/addons

Issue:
skip_gram_op fails when batch_size is provided
The python code for skip_gram_op calls "tf.train.batch", which doesn't exist for TF2 (https://github.com/tensorflow/addons/blob/master/tensorflow_addons/text/skip_gram_ops.py#L204). Note that the current tests don't test the batch_size argument, so is missing this scenario.



Discussion:
bot-of-gabrieldemarmiesse: @rahulunair 

You are owner of some files modified in this pull request.
Would you kindly review the changes whenever you have the time to?
Thank you very much.


Patch:
diff --git a/tensorflow_addons/layers/spectral_normalization.py b/tensorflow_addons/layers/spectral_normalization.py
index 03c392a..6990ffb 100644
--- a/tensorflow_addons/layers/spectral_normalization.py
+++ b/tensorflow_addons/layers/spectral_normalization.py
@@ -103,7 +103,6 @@ class SpectralNormalization(tf.keras.layers.Wrapper):
     def compute_output_shape(self, input_shape):
         return tf.TensorShape(self.layer.compute_output_shape(input_shape).as_list())
 
-    @tf.function
     def normalize_weights(self):
         """Generate spectral normalized weights.
 
@@ -118,11 +117,13 @@ class SpectralNormalization(tf.keras.layers.Wrapper):
             for _ in range(self.power_iterations):
                 v = tf.math.l2_normalize(tf.matmul(u, w, transpose_b=True))
                 u = tf.math.l2_normalize(tf.matmul(v, w))
-
+            u = tf.stop_gradient(u)
+            v = tf.stop_gradient(v)
             sigma = tf.matmul(tf.matmul(v, w), u, transpose_b=True)
-
-            self.w.assign(self.w / sigma)
-            self.u.assign(u)
+            self.u.assign(tf.cast(u, self.u.dtype))
+            self.w.assign(
+                tf.cast(tf.reshape(self.w / sigma, self.w_shape), self.w.dtype)
+            )
 
     def get_config(self):
         config = {"power_iterations": self.power_iterations}
diff --git a/tensorflow_addons/text/crf.py b/tensorflow_addons/text/crf.py
index e914f1d..e25cafe 100644
--- a/tensorflow_addons/text/crf.py
+++ b/tensorflow_addons/text/crf.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 import warnings
-from distutils.version import LooseVersion
 
 import numpy as np
 import tensorflow as tf
@@ -541,10 +540,6 @@ def crf_decode(
             "CRF Decoding does not work with KerasTensors in TF2.4. The bug has since been fixed in tensorflow/tensorflow##45534"
         )
 
-    if LooseVersion(tf.__version__) >= LooseVersion("2.5.0"):
-        warnings.warn(
-            "CRF decoding models have serialization issues in TF >=2.5 . Please see isse #2476"
-        )
     sequence_length = tf.cast(sequence_length, dtype=tf.int32)
 
     # If max_seq_len is 1, we skip the algorithm and simply return the
diff --git a/tensorflow_addons/text/skip_gram_ops.py b/tensorflow_addons/text/skip_gram_ops.py
index 51116d4..cd9f02b 100644
--- a/tensorflow_addons/text/skip_gram_ops.py
+++ b/tensorflow_addons/text/skip_gram_ops.py
@@ -38,8 +38,6 @@ def skip_gram_sample(
     vocab_min_count: Optional[FloatTensorLike] = None,
     vocab_subsampling: Optional[FloatTensorLike] = None,
     corpus_size: Optional[FloatTensorLike] = None,
-    batch_size: Optional[FloatTensorLike] = None,
-    batch_capacity: Optional[FloatTensorLike] = None,
     seed: Optional[FloatTensorLike] = None,
     name: Optional[str] = None,
 ) -> tf.Tensor:
@@ -76,10 +74,6 @@ def skip_gram_sample(
     down-sampled. See Eq. 5 in http://arxiv.org/abs/1310.4546 for more details
     about subsampling.
 
-    Due to the random window sizes used for each token, the lengths of the
-    outputs are non-deterministic, unless `batch_size` is specified to batch
-    the outputs to always return `Tensors` of length `batch_size`.
-
     Args:
       input_tensor: A rank-1 `Tensor` from which to generate skip-gram
         candidates.
@@ -121,20 +115,13 @@ def skip_gram_sample(
         counts of `vocab_freq_table`). Used with `vocab_subsampling` for
         down-sampling frequently occurring tokens. If this is specified,
         `vocab_freq_table` and `vocab_subsampling` must also be specified.
-      batch_size: (Optional) `int` specifying batch size of returned `Tensors`.
-      batch_capacity: (Optional) `int` specifying batch capacity for the queue
-        used for batching returned `Tensors`. Only has an effect if
-        `batch_size` > 0. Defaults to 100 * `batch_size` if not specified.
       seed: (Optional) `int` used to create a random seed for window size and
         subsampling. See `set_random_seed` docs for behavior.
       name: (Optional) A `string` name or a name scope for the operations.
 
     Returns:
       A `tuple` containing (token, label) `Tensors`. Each output `Tensor` is of
-      rank-1 and has the same type as `input_tensor`. The `Tensors` will be of
-      length `batch_size`; if `batch_size` is not specified, they will be of
-      random length, though they will be in sync with each other as long as
-      they are evaluated together.
+      rank-1 and has the same type as `input_tensor`.
 
     Raises:
       ValueError: If `vocab_freq_table` is not provided, but `vocab_min_count`,
@@ -193,18 +180,6 @@ def skip_gram_sample(
         # figures out sentence boundaries, then calls
         # skip_gram_generate_candidates() on each sentence.
 
-        # Batches the (tokens, labels) outputs so that they will be of deterministic
-        # batch_size, to facilitate feeding them into the rest of the network.
-        if batch_size is not None and batch_size > 0:
-            batch_capacity = (
-                batch_capacity
-                if (batch_capacity is not None and batch_capacity > 0)
-                else 100 * batch_size
-            )
-            return tf.train.batch(
-                [tokens, labels], batch_size, capacity=batch_capacity, enqueue_many=True
-            )
-
         return tokens, labels
 
 
@@ -224,8 +199,6 @@ def skip_gram_sample_with_text_vocab(
     start: FloatTensorLike = 0,
     limit: FloatTensorLike = -1,
     emit_self_as_target: bool = False,
-    batch_size: Optional[FloatTensorLike] = None,
-    batch_capacity: Optional[FloatTensorLike] = None,
     seed: Optional[FloatTensorLike] = None,
     name: Optional[str] = None,
 ) -> tf.Tensor:
@@ -298,10 +271,6 @@ def skip_gram_sample_with_text_vocab(
         to use the rest of the `Tensor` after `start`.
       emit_self_as_target: `bool` or scalar `Tensor` specifying whether to emit
         each token as a label for itself.
-      batch_size: (Optional) `int` specifying batch size of returned `Tensors`.
-      batch_capacity: (Optional) `int` specifying batch capacity for the queue
-        used for batching returned `Tensors`. Only has an effect if
-        `batch_size` > 0. Defaults to 100 * `batch_size` if not specified.
       seed: (Optional) `int` used to create a random seed for window size and
         subsampling. See
         [`set_random_seed`](../../g3doc/python/constant_op.md#set_random_seed)
@@ -310,10 +279,7 @@ def skip_gram_sample_with_text_vocab(
 
     Returns:
       A `tuple` containing (token, label) `Tensors`. Each output `Tensor` is of
-      rank-1 and has the same type as `input_tensor`. The `Tensors` will be of
-      length `batch_size`; if `batch_size` is not specified, they will be of
-      random length, though they will be in sync with each other as long as
-      they are evaluated together.
+      rank-1 and has the same type as `input_tensor`.
 
     Raises:
       ValueError: If `vocab_token_index` or `vocab_freq_index` is less than 0
@@ -397,8 +363,6 @@ def skip_gram_sample_with_text_vocab(
         vocab_subsampling=vocab_subsampling,
         # corpus_size is not used unless vocab_subsampling is specified.
         corpus_size=None if vocab_subsampling is None else corpus_size,
-        batch_size=batch_size,
-        batch_capacity=batch_capacity,
         seed=seed,
         name=name,
     )


Test patch:
diff --git a/tensorflow_addons/layers/tests/spectral_normalization_test.py b/tensorflow_addons/layers/tests/spectral_normalization_test.py
index a8aebab..b3f403c 100644
--- a/tensorflow_addons/layers/tests/spectral_normalization_test.py
+++ b/tensorflow_addons/layers/tests/spectral_normalization_test.py
@@ -109,6 +109,7 @@ def test_model_build(base_layer_fn, input_shape):
 
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
+@pytest.mark.usefixtures("run_with_mixed_precision_policy")
 def test_normalization():
     inputs = tf.keras.layers.Input(shape=[2, 2, 1])
 
diff --git a/tensorflow_addons/text/tests/crf_test.py b/tensorflow_addons/text/tests/crf_test.py
index c1da1c7..854c9fb 100644
--- a/tensorflow_addons/text/tests/crf_test.py
+++ b/tensorflow_addons/text/tests/crf_test.py
@@ -498,14 +498,18 @@ def test_tf_function():
     crf_decode(potentials, transition_params, sequence_length)
 
 
+class CRFDecode(tf.keras.layers.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, potentials, transition_params, sequence_length):
+        return text.crf_decode(potentials, transition_params, sequence_length)
+
+
 @pytest.mark.skipif(
     tf.__version__[:3] == "2.4",
     reason="CRF Decode doesn't work in TF2.4, the issue was fixed in TF core, but didn't make the release",
 )
-@pytest.mark.skipif(
-    tf.__version__[:3] >= "2.5",
-    reason="CRF decoding models have serialization issues in TF >=2.5 . Please see isse #2476",
-)
 def test_crf_decode_save_load(tmpdir):
     tf.keras.backend.clear_session()
     input_tensor = tf.keras.Input(shape=(10, 3), dtype=tf.float32, name="input_tensor")
@@ -513,7 +517,7 @@ def test_crf_decode_save_load(tmpdir):
     transition = tf.constant([[1, 1, 0], [0, 1, 1], [1, 0, 1]], dtype=tf.float32)
 
     output = tf.multiply(input_tensor, tf.constant(1.0))
-    decoded, _ = text.crf_decode(input_tensor, transition, seq_len)
+    decoded, _ = CRFDecode()(input_tensor, transition, seq_len)
 
     model = tf.keras.Model(
         inputs=[input_tensor, seq_len], outputs=[output, decoded], name="example_model"
@@ -541,7 +545,7 @@ def test_crf_decode_save_load(tmpdir):
     )
 
     temp_dir = str(tmpdir.mkdir("model"))
-    tf.saved_model.save(model, temp_dir)
+    model.save(temp_dir)
 
     tf.keras.backend.clear_session()
     model = tf.keras.models.load_model(