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:
Thanks, Can you send a small PR?
Looks like it would be slightly more involved than simply changing the symbol. I'm not super familiar with the testing framework, but it seems the tests are set up in eager mode (https://github.com/tensorflow/addons/blob/master/tensorflow_addons/text/tests/skip_gram_ops_test.py). If the code wants to support batch_size (ie. tf.compat.v1.train.batch), it seems the tests need to be rewritten in graph mode.  Any thoughts and pointers ?

Resolving patch:
diff --git a/tensorflow_addons/text/skip_gram_ops.py b/tensorflow_addons/text/skip_gram_ops.py
--- 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: