Merge of PR #1444

PiperOrigin-RevId: 233163674

Author: lgeiger

tensor2tensor/data_generators/problem.py

@@ -950,7 +950,7 @@ def _reverse_problem_hparams(p_hparams):
   # 'target', and each intended feature to swap has feature name 'input'.
   # In the future, remove need for this behavior.
   reversed_modality = {}
-  for feature_name in six.iterkeys(p.modality):
+  for feature_name in p.modality:
     reversed_feature_name = feature_name.replace("target", "input")
     if "target" in feature_name and reversed_feature_name in p.modality:
       reversed_modality[feature_name] = p.modality[reversed_feature_name]
@@ -962,7 +962,7 @@ def _reverse_problem_hparams(p_hparams):
 
   # Swap vocab sizes.
   reversed_vocab_size = {}
-  for feature_name in six.iterkeys(p.vocab_size):
+  for feature_name in p.vocab_size:
     reversed_feature_name = feature_name.replace("target", "input")
     if "target" in feature_name and reversed_feature_name in p.vocab_size:
       reversed_vocab_size[feature_name] = p.vocab_size[reversed_feature_name]

tensor2tensor/layers/latent_layers_test.py

@@ -20,7 +20,6 @@
 from __future__ import print_function
 
 import functools
-import six
 
 from tensor2tensor.layers import common_image_attention as cia
 from tensor2tensor.layers import discretization
@@ -138,8 +137,7 @@ def testTransformerAutoencoder(self):
     decoder_output, losses, cache = latent_layers.transformer_autoencoder(
         inputs, targets, target_space_id, hparams)
 
-    self.assertEqual(set(six.iterkeys(losses)),
-                     {"extra", "extra_loss", "latent_pred"})
+    self.assertEqual(set(losses), {"extra", "extra_loss", "latent_pred"})
 
     self.evaluate(tf.global_variables_initializer())
     decoder_output_, extra_loss_, latent_pred_ = self.evaluate(
@@ -154,5 +152,6 @@ def testTransformerAutoencoder(self):
     self.assertAllGreaterEqual(latent_pred_, 0.)
     self.assertEqual(cache, None)
 
+
 if __name__ == "__main__":
   tf.test.main()