internal merge of PR #999

PiperOrigin-RevId: 228622817

Author: gcampax

tensor2tensor/models/transformer.py

@@ -57,7 +57,7 @@ class Transformer(t2t_model.T2TModel):
 
   def __init__(self, *args, **kwargs):
     super(Transformer, self).__init__(*args, **kwargs)
-    self.attention_weights = dict()  # For visualizing attention heads.
+    self.attention_weights = {}  # For visualizing attention heads.
 
   def encode(self, inputs, target_space, hparams, features=None, losses=None):
     """Encode transformer inputs.
@@ -824,7 +824,7 @@ def fast_decode_tpu(encoder_output,
       hparams=hparams)
   if beam_size > 1:  # Beam Search
     initial_ids = sos_id * tf.ones([batch_size], dtype=tf.int32)
-    decoded_ids, scores = beam_search.beam_search(
+    decoded_ids, scores, _ = beam_search.beam_search(
         symbols_to_logits_fn,
         initial_ids,
         beam_size,
@@ -936,6 +936,7 @@ def fast_decode(encoder_output,
     force_decode_length: bool, whether to force the full decode length, or if
       False, stop when all beams hit eos_id.
     scope_prefix: str, prefix for decoder layer variable scopes.
+    cache: cache dictionary for additional predictions.
 
   Returns:
       A dict of decoding results {
@@ -959,7 +960,7 @@ def fast_decode(encoder_output,
       hparams.num_heads if hparams.get("attention_variables_3d") else 0)
 
   if cache is None:
-    cache = dict()
+    cache = {}
   cache.update({
       "layer_%d" % layer: {
           "k":

tensor2tensor/utils/beam_search.py

@@ -751,6 +751,13 @@ def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
     return tf.logical_and(
         tf.less(i, decode_length), tf.logical_not(bound_is_met))
 
+  inner_shape = tf.TensorShape([None, None, None])
+  if use_tpu:
+    inner_shape = tf.TensorShape([batch_size, beam_size, decode_length + 1])
+  if use_tpu:
+    state_struc = nest.map_structure(lambda state: state.get_shape(), states)
+  else:
+    state_struc = nest.map_structure(get_state_shape_invariants, states)
   (_, alive_seq, alive_log_probs, finished_seq, finished_scores,
    finished_flags, states) = tf.while_loop(
        _is_finished,
@@ -760,16 +767,12 @@ def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
        ],
        shape_invariants=[
            tf.TensorShape([]),
-           (tf.TensorShape([batch_size, beam_size, decode_length + 1])
-            if use_tpu else tf.TensorShape([None, None, None])),
+           inner_shape,
            alive_log_probs.get_shape(),
-           (tf.TensorShape([batch_size, beam_size, decode_length + 1])
-            if use_tpu else tf.TensorShape([None, None, None])),
+           inner_shape,
            finished_scores.get_shape(),
            finished_flags.get_shape(),
-           (nest.map_structure(lambda state: state.get_shape(), states)
-            if use_tpu else
-            nest.map_structure(get_state_shape_invariants, states)),
+           state_struc
        ],
        parallel_iterations=1,
        back_prop=False)

tensor2tensor/utils/beam_search_test.py

@@ -182,7 +182,7 @@ def symbols_to_logits(ids):
       logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
       return logits
 
-    final_ids, final_probs = beam_search.beam_search(
+    final_ids, final_probs, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,
@@ -390,10 +390,10 @@ def symbols_to_logits(ids, _, states):
         0.0,
         eos_id=1,
         states=states)
-    
+
     with self.test_session() as sess:
       final_states = sess.run(final_states)
-    self.assertAllEqual([[1]], final_states["state"])
+    self.assertAllEqual([[[2]]], final_states["state"])
 
   def testStateBeamTwo(self):
     batch_size = 1
@@ -476,7 +476,7 @@ def symbols_to_logits(_, i, states):
     states["state"] = tf.placeholder_with_default(
         states["state"], shape=(None, 1))
 
-    final_ids, _ = beam_search.beam_search(
+    final_ids, _, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,