Commit 005b8aa
committed
Fixes an issue with calling tf.contrib.seq2seq.dynamic_decode with an extended BasicDecoder which for example returns a tf.contrib.seq2seq.AttentionWrapperState.
In this case the internal while-loop fails when trying to store an instance tf.contrib.seq2seq.AttentionWrapperState in the internal TensorArray.
PiperOrigin-RevId: 1904917871 parent 9d8779e commit 005b8aa
2 files changed
Lines changed: 11 additions & 19 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -30,6 +30,7 @@ | |||
| 30 | 30 | from tensorflow.python.ops import control_flow_ops | |
| 31 | 31 | from tensorflow.python.ops import math_ops | |
| 32 | 32 | from tensorflow.python.ops import rnn | |
| 33 | + from tensorflow.python.ops import rnn_cell_impl | ||
| 33 | 34 | from tensorflow.python.ops import tensor_array_ops | |
| 34 | 35 | from tensorflow.python.ops import variable_scope | |
| 35 | 36 | from tensorflow.python.util import nest | |
@@ -39,6 +40,7 @@ | |||
| 39 | 40 | ||
| 40 | 41 | ||
| 41 | 42 | _transpose_batch_time = rnn._transpose_batch_time # pylint: disable=protected-access | |
| 43 | + _zero_state_tensors = rnn_cell_impl._zero_state_tensors # pylint: disable=protected-access | ||
| 42 | 44 | ||
| 43 | 45 | ||
| 44 | 46 | @six.add_metaclass(abc.ABCMeta) | |
@@ -133,16 +135,8 @@ def tracks_own_finished(self): | |||
| 133 | 135 | ||
| 134 | 136 | def _create_zero_outputs(size, dtype, batch_size): | |
| 135 | 137 | """Create a zero outputs Tensor structure.""" | |
| 136 | - def _t(s): | ||
| 137 | - return (s if isinstance(s, ops.Tensor) else constant_op.constant( | ||
| 138 | - tensor_shape.TensorShape(s).as_list(), | ||
| 139 | - dtype=dtypes.int32, | ||
| 140 | - name="zero_suffix_shape")) | ||
| 141 | - | ||
| 142 | 138 | def _create(s, d): | |
| 143 | - return array_ops.zeros( | ||
| 144 | - array_ops.concat( | ||
| 145 | - ([batch_size], _t(s)), axis=0), dtype=d) | ||
| 139 | + return _zero_state_tensors(s, batch_size, d) | ||
| 146 | 140 | ||
| 147 | 141 | return nest.map_structure(_create, size, dtype) | |
| 148 | 142 | ||
@@ -212,7 +206,8 @@ def dynamic_decode(decoder, | |||
| 212 | 206 | initial_time = constant_op.constant(0, dtype=dtypes.int32) | |
| 213 | 207 | ||
| 214 | 208 | def _shape(batch_size, from_shape): | |
| 215 | - if not isinstance(from_shape, tensor_shape.TensorShape): | ||
| 209 | + if (not isinstance(from_shape, tensor_shape.TensorShape) or | ||
| 210 | + from_shape.ndims == 0): | ||
| 216 | 211 | return tensor_shape.TensorShape(None) | |
| 217 | 212 | else: | |
| 218 | 213 | batch_size = tensor_util.constant_value( | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -49,24 +49,21 @@ | |||
| 49 | 49 | ||
| 50 | 50 | ||
| 51 | 51 | def _transpose_batch_time(x): | |
| 52 | - """Transpose the batch and time dimensions of a Tensor. | ||
| 52 | + """Transposes the batch and time dimensions of a Tensor. | ||
| 53 | 53 | ||
| 54 | - Retains as much of the static shape information as possible. | ||
| 54 | + If the input tensor has rank < 2 it returns the original tensor. Retains as | ||
| 55 | + much of the static shape information as possible. | ||
| 55 | 56 | ||
| 56 | 57 | Args: | |
| 57 | - x: A tensor of rank 2 or higher. | ||
| 58 | + x: A Tensor. | ||
| 58 | 59 | ||
| 59 | 60 | Returns: | |
| 60 | 61 | x transposed along the first two dimensions. | |
| 61 | - | ||
| 62 | - Raises: | ||
| 63 | - ValueError: if `x` is rank 1 or lower. | ||
| 64 | 62 | """ | |
| 65 | 63 | x_static_shape = x.get_shape() | |
| 66 | 64 | if x_static_shape.ndims is not None and x_static_shape.ndims < 2: | |
| 67 | - raise ValueError( | ||
| 68 | - "Expected input tensor %s to have rank at least 2, but saw shape: %s" % | ||
| 69 | - (x, x_static_shape)) | ||
| 65 | + return x | ||
| 66 | + | ||
| 70 | 67 | x_rank = array_ops.rank(x) | |
| 71 | 68 | x_t = array_ops.transpose( | |
| 72 | 69 | x, array_ops.concat( | |
0 commit comments