Add tpu saved model tags. No cpu tag is added because cpu is assumed to be the implicit device.

PiperOrigin-RevId: 176544698

1 files changed, 9 insertions, 1 deletions

diff --git a/tensorflow/contrib/saved_model/python/saved_model/reader_test.py b/tensorflow/contrib/saved_model/python/saved_model/reader_test.py
index a8331cbc8f..d10ec9cf0c 100644
--- a/tensorflow/contrib/saved_model/python/saved_model/reader_test.py
+++ b/tensorflow/contrib/saved_model/python/saved_model/reader_test.py
@@ -88,6 +88,13 @@ class ReaderTest(test.TestCase):
 
     # Graph that updates the single variable. SavedModel is invoked:
     # - to add the model (weights are not updated).
+    # - multiple predefined tags for serving on TPU.
+    with self.test_session(graph=ops.Graph()) as sess:
+      self._init_and_validate_variable(sess, "v", 44)
+      builder.add_meta_graph([tag_constants.SERVING, tag_constants.TPU])
+
+    # Graph that updates the single variable. SavedModel is invoked:
+    # - to add the model (weights are not updated).
     # - multiple custom tags.
     with self.test_session(graph=ops.Graph()) as sess:
       self._init_and_validate_variable(sess, "v", 45)
@@ -97,7 +104,8 @@ class ReaderTest(test.TestCase):
     builder.save()
 
     actual_tags = reader.get_saved_model_tag_sets(saved_model_dir)
-    expected_tags = [["train"], ["serve"], ["serve", "gpu"], ["foo", "bar"]]
+    expected_tags = [["train"], ["serve"], ["serve", "gpu"], ["serve", "tpu"],
+                     ["foo", "bar"]]
     self.assertEqual(expected_tags, actual_tags)