Merge of PR #1726

PiperOrigin-RevId: 281806525

Author: Seppo Enarvi

tensor2tensor/models/research/universal_transformer.py

@@ -458,7 +458,6 @@ def universal_transformer_base():
 @registry.register_hparams
 def universal_transformer_base_tpu():
   hparams = universal_transformer_base()
-  hparams = update_hparams_for_universal_transformer(hparams)
   transformer.update_hparams_for_tpu(hparams)
   hparams.add_step_timing_signal = False
   return hparams
@@ -467,7 +466,6 @@ def universal_transformer_base_tpu():
 @registry.register_hparams
 def universal_transformer_big():
   hparams = universal_transformer_base()
-  hparams = update_hparams_for_universal_transformer(hparams)
   hparams.hidden_size = 2048
   hparams.filter_size = 8192
   return hparams

tensor2tensor/models/transformer.py

@@ -1170,9 +1170,6 @@ def fast_decode(encoder_output,
           "scores": decoding log probs from the beam search,
               None if using greedy decoding (beam_size=1)
       }
-
-    Raises:
-      NotImplementedError: If beam size > 1 with partial targets.
   """
   if encoder_output is not None:
     batch_size = common_layers.shape_list(encoder_output)[0]

tensor2tensor/utils/decoding.py

@@ -927,6 +927,13 @@ def _interactive_input_tensor_to_features_dict(feature_map, hparams):
   features["decode_length"] = (
       IMAGE_DECODE_LENGTH if input_is_image else inputs[1])
   features["inputs"] = x
+  # Save inputs to "partial_targets" when prepending inputs to targets. Also
+  # keep "inputs" as some models crash if they don't exist.
+  if getattr(hparams, "prepend_mode", "none") != "none":
+    shape = tf.shape(x)
+    partial_targets = tf.reshape(x, [shape[0], shape[1]])
+    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
+    features["partial_targets"] = partial_targets
   return features
 
 
@@ -957,6 +964,13 @@ def _decode_input_tensor_to_features_dict(feature_map, hparams):
   features["decode_length"] = (
       IMAGE_DECODE_LENGTH if input_is_image else tf.shape(x)[1] + 50)
   features["inputs"] = x
+  # Save inputs to "partial_targets" when prepending inputs to targets. Also
+  # keep "inputs" as some models crash if they don't exist.
+  if getattr(hparams, "prepend_mode", "none") != "none":
+    shape = tf.shape(x)
+    partial_targets = tf.reshape(x, [shape[0], shape[1]])
+    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
+    features["partial_targets"] = partial_targets
   return features