Repository: tensorflow/addons

Issue:
Beam search decoder _check_static_batch_beam_maybe check correct?
**System information**
- TensorFlow installed from (source or binary): binary
- TensorFlow version (use command below): 2.0.0=alpha
- TensorFlow Addons installed from (source, PyPi): PyPi
- TensorFlow Addons version: 0.3.1
- Python version and type (eg. Anaconda Python, Stock Python as in Mac, or homebrew installed Python etc): Google Colab Python
- Is GPU used? (yes/no): yes
- GPU model (if used): T4

The check in beam_search_decoder.py#L183 seems faulty. As a MWE, consider:
```
tfa.seq2seq.beam_search_decoder._check_static_batch_beam_maybe(tf.TensorShape([3, 5, None]), 3, 5)
```
which yields:
```
TensorArray reordering expects elements to be reshapable to (3, 5, None) which is incompatible with the current shape (3, 5, None).
```

Either this error message is misleading or the check is faulty. Sadly, since I do not quite understand what it is supposed to be checking, I cannot propose a solution (other than to remove the last condition).


Discussion:
cc: @qlzh727 
I'm the original author of this function as part of https://github.com/tensorflow/tensorflow/commit/838a8f54f92452a15e3bb62a23ad5cd67e86933f and indeed the check is incorrect.

The check is simply: when the batch size and beam size are known statically, the leading dimensions should either be [batch_size*beam_size, ...] or [batch_size, beam_size, ...]. There is probably a cleaner way to do this.
Alright, cool! I'll try to submit a PR to fix make it do just that one of these days then :)
Ack, I can reproduce the error. @sjmielke feel free to send a PR, or you can assign to me if don't have time to fix it.

Resolving patch:
diff --git a/tensorflow_addons/seq2seq/beam_search_decoder.py b/tensorflow_addons/seq2seq/beam_search_decoder.py
--- a/tensorflow_addons/seq2seq/beam_search_decoder.py
+++ b/tensorflow_addons/seq2seq/beam_search_decoder.py
@@ -180,18 +180,22 @@ def _check_static_batch_beam_maybe(shape, batch_size, beam_width):
     """Raises an exception if dimensions are known statically and can not be
     reshaped to [batch_size, beam_size, -1]."""
     reshaped_shape = tf.TensorShape([batch_size, beam_width, None])
-    if (batch_size is not None and shape.dims[0].value is not None
-            and (shape[0] != batch_size * beam_width or
-                 (shape.ndims >= 2 and shape.dims[1].value is not None and
-                  (shape[0] != batch_size or shape[1] != beam_width)))):
-        tf.get_logger().warn(
-            "TensorArray reordering expects elements to be "
-            "reshapable to %s which is incompatible with the "
-            "current shape %s. Consider setting "
-            "reorder_tensor_arrays to False to disable TensorArray "
-            "reordering during the beam search." % (reshaped_shape, shape))
-        return False
-    return True
+    assert len(shape.dims) > 0
+    if batch_size is None or shape.dims[0].value is None:
+        return True  # not statically known => no check
+    if shape[0] == batch_size * beam_width:
+        return True  # flattened, matching
+    has_second_dim = shape.ndims >= 2 and shape.dims[1].value is not None
+    if has_second_dim and shape[0] == batch_size and shape[1] == beam_width:
+        return True  # non-flattened, matching
+    # Otherwise we could not find a match and warn:
+    tf.get_logger().warn(
+        "TensorArray reordering expects elements to be "
+        "reshapable to %s which is incompatible with the "
+        "current shape %s. Consider setting "
+        "reorder_tensor_arrays to False to disable TensorArray "
+        "reordering during the beam search." % (reshaped_shape, shape))
+    return False
 
 
 def _check_batch_beam(t, batch_size, beam_width):


Test patch:
diff --git a/tensorflow_addons/seq2seq/beam_search_decoder_test.py b/tensorflow_addons/seq2seq/beam_search_decoder_test.py
--- a/tensorflow_addons/seq2seq/beam_search_decoder_test.py
+++ b/tensorflow_addons/seq2seq/beam_search_decoder_test.py
@@ -209,6 +209,26 @@ def test_array_shape_dynamic_checks(self):
                                               5,
                                               is_valid=False)
 
+    def test_array_shape_static_checks(self):
+        self.assertTrue(
+            beam_search_decoder._check_static_batch_beam_maybe(
+                tf.TensorShape([None, None, None]), 3, 5))
+        self.assertTrue(
+            beam_search_decoder._check_static_batch_beam_maybe(
+                tf.TensorShape([15, None, None]), 3, 5))
+        self.assertFalse(
+            beam_search_decoder._check_static_batch_beam_maybe(
+                tf.TensorShape([16, None, None]), 3, 5))
+        self.assertTrue(
+            beam_search_decoder._check_static_batch_beam_maybe(
+                tf.TensorShape([3, 5, None]), 3, 5))
+        self.assertFalse(
+            beam_search_decoder._check_static_batch_beam_maybe(
+                tf.TensorShape([3, 6, None]), 3, 5))
+        self.assertFalse(
+            beam_search_decoder._check_static_batch_beam_maybe(
+                tf.TensorShape([5, 3, None]), 3, 5))
+
 
 class TestEosMasking(tf.test.TestCase):
     """Tests EOS masking used in beam search."""