Add distribution strategy tests for Keras preprocessing layers.

PiperOrigin-RevId: 309027642
Change-Id: Id38df0a26c9e134034f9b8ceadc7f9b91c219158

tensorflow/python/keras/layers/preprocessing/BUILD

@@ -266,17 +266,19 @@ tf_py_test(
     ],
 )
 
-tpu_py_test(
-    name = "categorical_encoding_tpu_test",
-    srcs = ["categorical_encoding_tpu_test.py"],
-    disable_experimental = True,
+distribute_py_test(
+    name = "categorical_encoding_distribution_test",
+    srcs = ["categorical_encoding_distribution_test.py"],
+    main = "categorical_encoding_distribution_test.py",
     python_version = "PY3",
-    tags = ["no_oss"],
+    tags = [
+        "multi_and_single_gpu",
+    ],
     deps = [
         ":categorical_encoding",
-        "//tensorflow/python/distribute:tpu_strategy",
+        "//tensorflow/python/distribute:combinations",
+        "//tensorflow/python/distribute:strategy_combinations",
         "//tensorflow/python/keras",
-        "//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
     ],
 )
 
@@ -293,17 +295,17 @@ tf_py_test(
     ],
 )
 
-tpu_py_test(
-    name = "discretization_tpu_test",
-    srcs = ["discretization_tpu_test.py"],
-    disable_experimental = True,
+distribute_py_test(
+    name = "discretization_distribution_test",
+    srcs = ["discretization_distribution_test.py"],
+    main = "discretization_distribution_test.py",
     python_version = "PY3",
-    tags = ["no_oss"],
+    tags = ["multi_and_single_gpu"],
     deps = [
         ":discretization",
-        "//tensorflow/python/distribute:tpu_strategy",
+        "//tensorflow/python/distribute:combinations",
+        "//tensorflow/python/distribute:strategy_combinations",
         "//tensorflow/python/keras",
-        "//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
     ],
 )
 
@@ -322,16 +324,16 @@ cuda_py_test(
 )
 
 tpu_py_test(
-    name = "hashing_tpu_test",
-    srcs = ["hashing_tpu_test.py"],
-    disable_experimental = True,
+    name = "hashing_distribution_test",
+    srcs = ["hashing_distribution_test.py"],
+    main = "hashing_distribution_test.py",
     python_version = "PY3",
-    tags = ["no_oss"],
+    tags = ["multi_and_single_gpu"],
     deps = [
         ":hashing",
-        "//tensorflow/python/distribute:tpu_strategy",
+        "//tensorflow/python/distribute:combinations",
+        "//tensorflow/python/distribute:strategy_combinations",
         "//tensorflow/python/keras",
-        "//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
     ],
 )
 
@@ -352,16 +354,16 @@ tf_py_test(
 )
 
 tpu_py_test(
-    name = "index_lookup_tpu_test",
-    srcs = ["index_lookup_tpu_test.py"],
-    disable_experimental = True,
+    name = "index_lookup_distribution_test",
+    srcs = ["index_lookup_distribution_test.py"],
+    main = "index_lookup_distribution_test.py",
     python_version = "PY3",
     tags = ["no_oss"],
     deps = [
         ":index_lookup",
-        "//tensorflow/python/distribute:tpu_strategy",
+        "//tensorflow/python/distribute:combinations",
+        "//tensorflow/python/distribute:strategy_combinations",
         "//tensorflow/python/keras",
-        "//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
     ],
 )
 

tensorflow/python/keras/layers/preprocessing/categorical_encoding_distribution_test.py

@@ -21,21 +21,24 @@ from __future__ import print_function
 import numpy as np
 
 from tensorflow.python import keras
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import strategy_combinations
 from tensorflow.python.framework import dtypes
 from tensorflow.python.keras import keras_parameterized
-from tensorflow.python.keras.distribute import tpu_strategy_test_utils
 from tensorflow.python.keras.layers.preprocessing import categorical_encoding
 from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils
 from tensorflow.python.platform import test
 
 
-@keras_parameterized.run_all_keras_modes(
-    always_skip_v1=True, always_skip_eager=True)
+@combinations.generate(
+    combinations.combine(
+        distribution=strategy_combinations.all_strategies,
+        mode=["eager", "graph"]))
 class CategoricalEncodingDistributionTest(
     keras_parameterized.TestCase,
     preprocessing_test_utils.PreprocessingLayerTest):
 
-  def test_tpu_distribution(self):
+  def test_distribution(self, distribution):
     input_array = np.array([[1, 2, 3, 1], [0, 3, 1, 0]])
 
     # pyformat: disable
@@ -44,9 +47,7 @@ class CategoricalEncodingDistributionTest(
     # pyformat: enable
     max_tokens = 6
 
-    strategy = tpu_strategy_test_utils.get_tpu_strategy()
-
-    with strategy.scope():
+    with distribution.scope():
       input_data = keras.Input(shape=(4,), dtype=dtypes.int32)
       layer = categorical_encoding.CategoricalEncoding(
           max_tokens=max_tokens, output_mode=categorical_encoding.BINARY)

tensorflow/python/keras/layers/preprocessing/discretization_distribution_test.py

@@ -21,27 +21,29 @@ from __future__ import print_function
 import numpy as np
 
 from tensorflow.python import keras
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import strategy_combinations
 from tensorflow.python.keras import keras_parameterized
-from tensorflow.python.keras.distribute import tpu_strategy_test_utils
 from tensorflow.python.keras.layers.preprocessing import discretization
 from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils
 from tensorflow.python.platform import test
 
 
-@keras_parameterized.run_all_keras_modes(
-    always_skip_v1=True, always_skip_eager=True)
+@combinations.generate(
+    combinations.combine(
+        distribution=strategy_combinations.all_strategies,
+        mode=["eager", "graph"]))
 class DiscretizationDistributionTest(
     keras_parameterized.TestCase,
     preprocessing_test_utils.PreprocessingLayerTest):
 
-  def test_tpu_distribution(self):
+  def test_distribution(self, distribution):
     input_array = np.array([[-1.5, 1.0, 3.4, .5], [0.0, 3.0, 1.3, 0.0]])
 
     expected_output = [[0, 2, 3, 1], [1, 3, 2, 1]]
     expected_output_shape = [None, None]
 
-    strategy = tpu_strategy_test_utils.get_tpu_strategy()
-    with strategy.scope():
+    with distribution.scope():
       input_data = keras.Input(shape=(None,))
       layer = discretization.Discretization(
           bins=[0., 1., 2.], output_mode=discretization.INTEGER)

tensorflow/python/keras/layers/preprocessing/hashing_distribution_test.py

@@ -22,30 +22,32 @@ import numpy as np
 
 from tensorflow.python import keras
 from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import strategy_combinations
 from tensorflow.python.framework import config
 from tensorflow.python.framework import dtypes
 from tensorflow.python.keras import keras_parameterized
-from tensorflow.python.keras.distribute import tpu_strategy_test_utils
 from tensorflow.python.keras.layers.preprocessing import hashing
 from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils
 from tensorflow.python.platform import test
 
 
-@keras_parameterized.run_all_keras_modes(
-    always_skip_v1=True, always_skip_eager=True)
+@combinations.generate(
+    combinations.combine(
+        distribution=strategy_combinations.all_strategies,
+        mode=["eager", "graph"]))
 class HashingDistributionTest(keras_parameterized.TestCase,
                               preprocessing_test_utils.PreprocessingLayerTest):
 
-  def test_tpu_distribution(self):
+  def test_distribution(self, distribution):
     input_data = np.asarray([["omar"], ["stringer"], ["marlo"], ["wire"]])
     input_dataset = dataset_ops.Dataset.from_tensor_slices(input_data).batch(
         2, drop_remainder=True)
     expected_output = [[0], [0], [1], [0]]
 
     config.set_soft_device_placement(True)
-    strategy = tpu_strategy_test_utils.get_tpu_strategy()
 
-    with strategy.scope():
+    with distribution.scope():
       input_data = keras.Input(shape=(None,), dtype=dtypes.string)
       layer = hashing.Hashing(num_bins=2)
       int_data = layer(input_data)

tensorflow/python/keras/layers/preprocessing/index_lookup_distribution_test.py

@@ -22,22 +22,34 @@ import numpy as np
 
 from tensorflow.python import keras
 from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import strategy_combinations
+from tensorflow.python.eager import context
 from tensorflow.python.framework import config
 from tensorflow.python.framework import dtypes
 from tensorflow.python.keras import keras_parameterized
-from tensorflow.python.keras.distribute import tpu_strategy_test_utils
 from tensorflow.python.keras.layers.preprocessing import index_lookup
+from tensorflow.python.keras.layers.preprocessing import index_lookup_v1
 from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils
 from tensorflow.python.platform import test
 
 
-@keras_parameterized.run_all_keras_modes(
-    always_skip_v1=True, always_skip_eager=True)
+def get_layer_class():
+  if context.executing_eagerly():
+    return index_lookup.IndexLookup
+  else:
+    return index_lookup_v1.IndexLookup
+
+
+@combinations.generate(
+    combinations.combine(
+        distribution=strategy_combinations.all_strategies,
+        mode=["eager", "graph"]))
 class IndexLookupDistributionTest(
     keras_parameterized.TestCase,
     preprocessing_test_utils.PreprocessingLayerTest):
 
-  def test_tpu_distribution(self):
+  def test_tpu_distribution(self, distribution):
     vocab_data = [[
         "earth", "earth", "earth", "earth", "wind", "wind", "wind", "and",
         "and", "fire"
@@ -50,11 +62,10 @@ class IndexLookupDistributionTest(
     expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
 
     config.set_soft_device_placement(True)
-    strategy = tpu_strategy_test_utils.get_tpu_strategy()
 
-    with strategy.scope():
+    with distribution.scope():
       input_data = keras.Input(shape=(None,), dtype=dtypes.string)
-      layer = index_lookup.IndexLookup()
+      layer = get_layer_class()()
       layer.adapt(vocab_dataset)
       int_data = layer(input_data)
       model = keras.Model(inputs=input_data, outputs=int_data)