Make sure we run tf_upgrade_v2 only in v1 and testLinearClassifier after v2 conversion only with TF 2.x.

PiperOrigin-RevId: 260582754

tensorflow/tools/compatibility/BUILD

@@ -153,6 +153,7 @@ py_test(
     srcs = ["tf_upgrade_v2_test.py"],
     python_version = "PY2",
     srcs_version = "PY2AND3",
+    tags = ["v1only"],
     deps = [
         ":tf_upgrade_v2_lib",
         "//tensorflow:tensorflow_py",
@@ -225,7 +226,8 @@ genrule(
     cmd = ("$(location :tf_upgrade_v2)" +
            " --infile $(location testdata/test_file_v1_12.py)" +
            " --outfile $(location test_file_v2_0.py)" +
-           " --reportfile $(location report_v2.txt)"),
+           " --reportfile $(location report_v2.txt) && " +
+           "sed -i 's/_TEST_VERSION = 1/_TEST_VERSION = 2/g' $(location test_file_v2_0.py)"),
     tools = [":tf_upgrade_v2"],
 )
 
@@ -235,6 +237,7 @@ py_test(
     srcs = ["testdata/test_file_v1_12.py"],
     python_version = "PY2",
     srcs_version = "PY2AND3",
+    tags = ["v1only"],
     deps = [
         "//tensorflow:tensorflow_py",
     ],

tensorflow/tools/compatibility/testdata/test_file_v1_12.py

@@ -21,10 +21,16 @@ import tensorflow as tf
 from tensorflow.python.framework import test_util
 from tensorflow.python.platform import test as test_lib
 
+_TEST_VERSION = 1
+
 
 class TestUpgrade(test_util.TensorFlowTestCase):
   """Test various APIs that have been changed in 2.0."""
 
+  @classmethod
+  def setUpClass(cls):
+    cls._tf_api_version = 1 if hasattr(tf, 'contrib') else 2
+
   def setUp(self):
     tf.compat.v1.enable_v2_behavior()
 
@@ -74,6 +80,14 @@ class TestUpgrade(test_util.TensorFlowTestCase):
     self.assertAllClose(out, 0.40318608)
 
   def testLinearClassifier(self):
+    if _TEST_VERSION == 2 and self._tf_api_version == 1:
+      # Skip if we converted this file to v2 but running with tf v1.
+      # In this case, conversion script adds reference to
+      # tf.keras.losses.Reduction which is not available in v1.
+      self.skipTest(
+          'After converting to 2.0, this test does not work with '
+          'TensorFlow 1.x.')
+      return
     feature_column = tf.feature_column.numeric_column(
         'feature', shape=(1,))
 

tensorflow/tools/compatibility/tf_upgrade_v2.py

@@ -2024,7 +2024,7 @@ def _add_loss_reduction_transformer(parent, node, full_name, name, logs):
 
   Default value for tf.estimator.*Classifier and tf.estimator.*Regressor
   loss_reduction argument changed to SUM_OVER_BATCH_SIZE. So, we update
-  existing calls to use the old default value `tf.losses.Reduction.SUM`.
+  existing calls to use the old default value `tf.keras.losses.Reduction.SUM`.
 
   Note: to apply this transformation, symbol must be added
   to reordered_function_names above.
@@ -2032,9 +2032,7 @@ def _add_loss_reduction_transformer(parent, node, full_name, name, logs):
   for keyword_arg in node.keywords:
     if keyword_arg.arg == "loss_reduction":
       return node
-  # TODO(annarev): this should be updated to tf.keras.losses.Reduction.SUM
-  # once b/125525822 is fixed.
-  default_value = "tf.compat.v1.losses.Reduction.SUM"
+  default_value = "tf.keras.losses.Reduction.SUM"
   # Parse with pasta instead of ast to avoid emitting a spurious trailing \n.
   ast_value = pasta.parse(default_value)
   node.keywords.append(ast.keyword(arg="loss_reduction", value=ast_value))

tensorflow/tools/compatibility/tf_upgrade_v2_test.py

@@ -684,7 +684,7 @@ bazel-bin/tensorflow/tools/compatibility/update/generate_v2_reorders_map
     for c in classes:
       ns = "tf.estimator." + c
       text = ns + "()"
-      expected_text = ns + "(loss_reduction=tf.compat.v1.losses.Reduction.SUM)"
+      expected_text = ns + "(loss_reduction=tf.keras.losses.Reduction.SUM)"
       _, report, errors, new_text = self._upgrade(text)
       self.assertEqual(expected_text, new_text)
 
@@ -703,7 +703,7 @@ bazel-bin/tensorflow/tools/compatibility/update/generate_v2_reorders_map
     text = "tf.estimator.BaselineClassifier(model_dir=model_dir)"
     expected_text = ("tf.estimator.BaselineClassifier(" +
                      "model_dir=model_dir, "
-                     "loss_reduction=tf.compat.v1.losses.Reduction.SUM)")
+                     "loss_reduction=tf.keras.losses.Reduction.SUM)")
     _, report, errors, new_text = self._upgrade(text)
     self.assertEqual(expected_text, new_text)
 
@@ -728,7 +728,7 @@ bazel-bin/tensorflow/tools/compatibility/update/generate_v2_reorders_map
       suffix = "(input_layer_partitioner=TEST)"
       text = ns + suffix
       suffix = ("(input_layer_partitioner=TEST, "
-                "loss_reduction=tf.compat.v1.losses.Reduction.SUM)")
+                "loss_reduction=tf.keras.losses.Reduction.SUM)")
       expected_text = "tf.compat.v1.estimator." + c + suffix
       _, unused_report, unused_errors, new_text = self._upgrade(text)
       self.assertEqual(new_text, expected_text)
@@ -764,7 +764,7 @@ bazel-bin/tensorflow/tools/compatibility/update/generate_v2_reorders_map
       suffix = "(optimizer=TEST)"
       text = ns + suffix
       suffix = ("(optimizer=TEST, "
-                "loss_reduction=tf.compat.v1.losses.Reduction.SUM)")
+                "loss_reduction=tf.keras.losses.Reduction.SUM)")
       expected_text = "tf.compat.v1.estimator." + c + suffix
       _, unused_report, unused_errors, new_text = self._upgrade(text)
       self.assertEqual(new_text, expected_text)
@@ -779,7 +779,7 @@ bazel-bin/tensorflow/tools/compatibility/update/generate_v2_reorders_map
       suffix = "(dnn_optimizer=TEST, linear_optimizer=Test)"
       text = ns + suffix
       suffix = ("(dnn_optimizer=TEST, linear_optimizer=Test, "
-                "loss_reduction=tf.compat.v1.losses.Reduction.SUM)")
+                "loss_reduction=tf.keras.losses.Reduction.SUM)")
       expected_text = "tf.compat.v1.estimator." + c + suffix
       _, unused_report, unused_errors, new_text = self._upgrade(text)
       self.assertEqual(new_text, expected_text)
@@ -815,7 +815,7 @@ bazel-bin/tensorflow/tools/compatibility/update/generate_v2_reorders_map
       suffix = "(input_layer_partitioner=TEST, optimizer=TEST)"
       text = ns + suffix
       suffix = ("(input_layer_partitioner=TEST, optimizer=TEST, "
-                "loss_reduction=tf.compat.v1.losses.Reduction.SUM)")
+                "loss_reduction=tf.keras.losses.Reduction.SUM)")
       expected_text = "tf.compat.v1.estimator." + c + suffix
       _, unused_report, unused_errors, new_text = self._upgrade(text)
       self.assertEqual(new_text, expected_text)
@@ -833,7 +833,7 @@ bazel-bin/tensorflow/tools/compatibility/update/generate_v2_reorders_map
       text = ns + suffix
       suffix = ("(input_layer_partitioner=TEST, dnn_optimizer=TEST, "
                 "linear_optimizer=TEST, "
-                "loss_reduction=tf.compat.v1.losses.Reduction.SUM)")
+                "loss_reduction=tf.keras.losses.Reduction.SUM)")
       expected_text = "tf.compat.v1.estimator." + c + suffix
       _, unused_report, unused_errors, new_text = self._upgrade(text)
       self.assertEqual(new_text, expected_text)