Update RNN helpers to be able to handle dynamic state sizes.

This fixes a bug I introduced previously by adding the alignment into the AttentionWrapper's state (since the alignment's size may have to be a Tensor - the encoder's max_time is not usually static).

PiperOrigin-RevId: 156077314

Author: ebrevdo

tensorflow/contrib/learn/python/learn/estimators/state_saving_rnn_estimator_test.py

@@ -525,7 +525,7 @@ def testLearnShiftByOne(self):
     num_classes = 2
     num_unroll = 32
     sequence_length = 32
-    train_steps = 200
+    train_steps = 300
     eval_steps = 20
     num_units = [4]
     learning_rate = 0.5

tensorflow/contrib/legacy_seq2seq/python/kernel_tests/seq2seq_test.py

@@ -942,8 +942,8 @@ def SampledLoss(labels, logits):
         perplexities[bucket].append(math.exp(float(res[1])))
       for bucket in range(len(buckets)):
         if len(perplexities[bucket]) > 1:  # Assert that perplexity went down.
-          self.assertLess(perplexities[bucket][-1],  # 10% margin of error.
-                          1.1 * perplexities[bucket][0])
+          self.assertLess(perplexities[bucket][-1],  # 20% margin of error.
+                          1.2 * perplexities[bucket][0])
 
   def testModelWithBooleanFeedPrevious(self):
     """Test the model behavior when feed_previous is True.

tensorflow/contrib/rnn/python/ops/core_rnn.py

@@ -31,7 +31,7 @@
 
 
 # pylint: disable=protected-access
-_state_size_with_prefix = rnn_cell_impl._state_size_with_prefix
+_concat = rnn_cell_impl._concat
 _infer_state_dtype = rnn._infer_state_dtype
 _reverse_seq = rnn._reverse_seq
 _rnn_step = rnn._rnn_step
@@ -159,11 +159,10 @@ def static_rnn(cell, inputs, initial_state=None, dtype=None,
             "sequence_length must be a vector of length batch_size")
       def _create_zero_output(output_size):
         # convert int to TensorShape if necessary
-        size = _state_size_with_prefix(output_size, prefix=[batch_size])
+        size = _concat(batch_size, output_size)
         output = array_ops.zeros(
             array_ops.stack(size), _infer_state_dtype(dtype, state))
-        shape = _state_size_with_prefix(
-            output_size, prefix=[fixed_batch_size.value])
+        shape = _concat(fixed_batch_size.value, output_size, static=True)
         output.set_shape(tensor_shape.TensorShape(shape))
         return output
 

tensorflow/contrib/seq2seq/python/kernel_tests/basic_decoder_test.py

@@ -124,7 +124,7 @@ def testStepWithGreedyEmbeddingHelper(self):
     vocabulary_size = 7
     cell_depth = vocabulary_size  # cell's logits must match vocabulary size
     input_depth = 10
-    start_tokens = [0] * batch_size
+    start_tokens = np.random.randint(0, vocabulary_size, size=batch_size)
     end_token = 1
 
     with self.test_session(use_gpu=True) as sess:

tensorflow/python/ops/rnn.py

@@ -33,7 +33,7 @@
 
 
 # pylint: disable=protected-access
-_state_size_with_prefix = rnn_cell_impl._state_size_with_prefix
+_concat = rnn_cell_impl._concat
 # pylint: enable=protected-access
 
 
@@ -660,7 +660,7 @@ def _dynamic_rnn_loop(cell,
 
   # Prepare dynamic conditional copying of state & output
   def _create_zero_arrays(size):
-    size = _state_size_with_prefix(size, prefix=[batch_size])
+    size = _concat(batch_size, size)
     return array_ops.zeros(
         array_ops.stack(size), _infer_state_dtype(dtype, state))
 
@@ -746,8 +746,8 @@ def _time_step(time, output_ta_t, state):
 
   # Restore some shape information
   for output, output_size in zip(final_outputs, flat_output_size):
-    shape = _state_size_with_prefix(
-        output_size, prefix=[const_time_steps, const_batch_size])
+    shape = _concat(
+        [const_time_steps, const_batch_size], output_size, static=True)
     output.set_shape(shape)
 
   final_outputs = nest.pack_sequence_as(
@@ -981,9 +981,7 @@ def loop_fn(time, cell_output, cell_state, loop_state):
     emit_ta = nest.pack_sequence_as(structure=emit_structure,
                                     flat_sequence=flat_emit_ta)
     flat_zero_emit = [
-        array_ops.zeros(
-            _state_size_with_prefix(size_i, prefix=[batch_size]),
-            dtype_i)
+        array_ops.zeros(_concat(batch_size, size_i), dtype_i)
         for size_i, dtype_i in zip(flat_emit_size, flat_emit_dtypes)]
     zero_emit = nest.pack_sequence_as(structure=emit_structure,
                                       flat_sequence=flat_zero_emit)

tensorflow/python/ops/rnn_cell_impl.py

@@ -26,55 +26,82 @@
 
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_util
 from tensorflow.python.layers import base as base_layer
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import variable_scope as vs
 from tensorflow.python.ops import variables as tf_variables
 from tensorflow.python.util import nest
 
 
-def _state_size_with_prefix(state_size, prefix=None):
-  """Helper function that enables int or TensorShape shape specification.
+def _concat(prefix, suffix, static=False):
+  """Concat that enables int, Tensor, or TensorShape values.
 
-  This function takes a size specification, which can be an integer or a
-  TensorShape, and converts it into a list of integers. One may specify any
-  additional dimensions that precede the final state size specification.
+  This function takes a size specification, which can be an integer, a
+  TensorShape, or a Tensor, and converts it into a concatenated Tensor
+  (if static = False) or a list of integers (if static = True).
 
   Args:
-    state_size: TensorShape or int that specifies the size of a tensor.
-    prefix: optional additional list of dimensions to prepend.
+    prefix: The prefix; usually the batch size (and/or time step size).
+      (TensorShape, int, or Tensor.)
+    suffix: TensorShape, int, or Tensor.
+    static: If `True`, return a python list with possibly unknown dimensions.
+      Otherwise return a `Tensor`.
 
   Returns:
-    result_state_size: list of dimensions the resulting tensor size.
+    shape: the concatenation of prefix and suffix.
+
+  Raises:
+    ValueError: if `suffix` is not a scalar or vector (or TensorShape).
+    ValueError: if prefix or suffix was `None` and asked for dynamic
+      Tensors out.
   """
-  result_state_size = tensor_shape.as_shape(state_size).as_list()
-  if prefix is not None:
-    if not isinstance(prefix, list):
-      raise TypeError("prefix of _state_size_with_prefix should be a list.")
-    result_state_size = prefix + result_state_size
-  return result_state_size
+  if isinstance(prefix, ops.Tensor):
+    p = prefix
+    p_static = tensor_util.constant_value(prefix)
+    if p.shape.ndims == 0:
+      p = array_ops.expand_dims(p, 0)
+    elif p.shape.ndims != 1:
+      raise ValueError("prefix tensor must be either a scalar or vector, "
+                       "but saw tensor: %s" % p)
+  else:
+    p = tensor_shape.as_shape(prefix)
+    p = p.as_list() if p.ndims is not None else None
+    p_static = p
+  if isinstance(suffix, ops.Tensor):
+    s = suffix
+    s_static = tensor_util.constant_value(suffix)
+    if s.shape.ndims == 0:
+      s = array_ops.expand_dims(s, 0)
+    elif s.shape.ndims != 1:
+      raise ValueError("suffix tensor must be either a scalar or vector, "
+                       "but saw tensor: %s" % s)
+  else:
+    s = tensor_shape.as_shape(suffix)
+    s = s.as_list() if s.ndims is not None else None
+    s_static = s
+
+  if static:
+    shape = tensor_shape.as_shape(p_static).concatenate(s_static)
+    shape = shape.as_list() if shape.ndims is not None else None
+  else:
+    if p is None or s is None:
+      raise ValueError("Provided a prefix or suffix of None: %s and %s"
+                       % (prefix, suffix))
+    shape = array_ops.concat((p, s), 0)
+  return shape
 
 
 def _zero_state_tensors(state_size, batch_size, dtype):
   """Create tensors of zeros based on state_size, batch_size, and dtype."""
-  if nest.is_sequence(state_size):
-    state_size_flat = nest.flatten(state_size)
-    zeros_flat = [
-        array_ops.zeros(
-            array_ops.stack(_state_size_with_prefix(
-                s, prefix=[batch_size])),
-            dtype=dtype) for s in state_size_flat
-    ]
-    for s, z in zip(state_size_flat, zeros_flat):
-      z.set_shape(_state_size_with_prefix(s, prefix=[None]))
-    zeros = nest.pack_sequence_as(structure=state_size,
-                                  flat_sequence=zeros_flat)
-  else:
-    zeros_size = _state_size_with_prefix(state_size, prefix=[batch_size])
-    zeros = array_ops.zeros(array_ops.stack(zeros_size), dtype=dtype)
-    zeros.set_shape(_state_size_with_prefix(state_size, prefix=[None]))
-
-  return zeros
+  def get_state_shape(s):
+    """Combine s with batch_size to get a proper tensor shape."""
+    c = _concat(batch_size, s)
+    c_static = _concat(batch_size, s, static=True)
+    size = array_ops.zeros(c, dtype=dtype)
+    size.set_shape(c_static)
+    return size
+  return nest.map_structure(get_state_shape, state_size)
 
 
 class _RNNCell(base_layer.Layer):