Add distribution strategy tests.

PiperOrigin-RevId: 308090534
Change-Id: I70c08479f4ae324f3d7a8835bf89b0ac10105811

tensorflow/python/keras/distribute/BUILD

@@ -440,3 +440,15 @@ py_library(
         "//tensorflow/python/keras/optimizer_v2",
     ],
 )
+
+py_library(
+    name = "tpu_strategy_test_utils",
+    srcs = ["tpu_strategy_test_utils.py"],
+    deps = [
+        "//tensorflow/python:platform",
+        "//tensorflow/python/distribute:tpu_strategy",
+        "//tensorflow/python/distribute/cluster_resolver:tpu_cluster_resolver_py",
+        "//tensorflow/python/eager:remote",
+        "//tensorflow/python/tpu:tpu_strategy_util",
+    ],
+)

tensorflow/python/keras/distribute/tpu_strategy_test_utils.py

@@ -0,0 +1,47 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utility functions for tests using TPUStrategy."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import google_type_annotations
+from __future__ import print_function
+
+from tensorflow.python.distribute import tpu_strategy
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
+from tensorflow.python.eager import remote
+from tensorflow.python.platform import flags
+from tensorflow.python.tpu import tpu_strategy_util
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string("tpu", "", "Name of TPU to connect to.")
+flags.DEFINE_string("project", None, "Name of GCP project with TPU.")
+flags.DEFINE_string("zone", None, "Name of GCP zone with TPU.")
+
+
+def get_tpu_cluster_resolver():
+  resolver = tpu_cluster_resolver.TPUClusterResolver(
+      tpu=FLAGS.tpu,
+      zone=FLAGS.zone,
+      project=FLAGS.project,
+  )
+  return resolver
+
+
+def get_tpu_strategy():
+  resolver = get_tpu_cluster_resolver()
+  remote.connect_to_cluster(resolver)
+  tpu_strategy_util.initialize_tpu_system(resolver)
+  return tpu_strategy.TPUStrategy(resolver)

tensorflow/python/keras/layers/preprocessing/BUILD

@@ -3,6 +3,7 @@
 
 load("//tensorflow:tensorflow.bzl", "tf_py_test")
 load("//tensorflow:tensorflow.bzl", "cuda_py_test")
+load("//tensorflow/python/tpu:tpu.bzl", "tpu_py_test")
 
 package(
     default_visibility = [
@@ -231,19 +232,6 @@ py_library(
     ],
 )
 
-tf_py_test(
-    name = "discretization_test",
-    size = "small",
-    srcs = ["discretization_test.py"],
-    python_version = "PY3",
-    deps = [
-        ":discretization",
-        ":preprocessing_test_utils",
-        "//tensorflow/python:client_testlib",
-        "@absl_py//absl/testing:parameterized",
-    ],
-)
-
 cuda_py_test(
     name = "categorical_crossing_test",
     size = "medium",
@@ -261,6 +249,63 @@ cuda_py_test(
     ],
 )
 
+tf_py_test(
+    name = "categorical_encoding_test",
+    size = "medium",
+    srcs = ["categorical_encoding_test.py"],
+    python_version = "PY3",
+    deps = [
+        ":categorical_encoding",
+        ":preprocessing_test_utils",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python/keras",
+        "//tensorflow/python/keras/utils:generic_utils",
+        "//tensorflow/python/ops/ragged:ragged_string_ops",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_test(
+    name = "categorical_encoding_tpu_test",
+    srcs = ["categorical_encoding_tpu_test.py"],
+    disable_experimental = True,
+    python_version = "PY3",
+    tags = ["no_oss"],
+    deps = [
+        ":categorical_encoding",
+        "//tensorflow/python/distribute:tpu_strategy",
+        "//tensorflow/python/keras",
+        "//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
+    ],
+)
+
+tf_py_test(
+    name = "discretization_test",
+    size = "small",
+    srcs = ["discretization_test.py"],
+    python_version = "PY3",
+    deps = [
+        ":discretization",
+        ":preprocessing_test_utils",
+        "//tensorflow/python:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_test(
+    name = "discretization_tpu_test",
+    srcs = ["discretization_tpu_test.py"],
+    disable_experimental = True,
+    python_version = "PY3",
+    tags = ["no_oss"],
+    deps = [
+        ":discretization",
+        "//tensorflow/python/distribute:tpu_strategy",
+        "//tensorflow/python/keras",
+        "//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
+    ],
+)
+
 cuda_py_test(
     name = "hashing_test",
     size = "medium",
@@ -275,6 +320,20 @@ cuda_py_test(
     ],
 )
 
+tpu_py_test(
+    name = "hashing_tpu_test",
+    srcs = ["hashing_tpu_test.py"],
+    disable_experimental = True,
+    python_version = "PY3",
+    tags = ["no_oss"],
+    deps = [
+        ":hashing",
+        "//tensorflow/python/distribute:tpu_strategy",
+        "//tensorflow/python/keras",
+        "//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
+    ],
+)
+
 tf_py_test(
     name = "index_lookup_test",
     size = "medium",
@@ -291,6 +350,20 @@ tf_py_test(
     ],
 )
 
+tpu_py_test(
+    name = "index_lookup_tpu_test",
+    srcs = ["index_lookup_tpu_test.py"],
+    disable_experimental = True,
+    python_version = "PY3",
+    tags = ["no_oss"],
+    deps = [
+        ":index_lookup",
+        "//tensorflow/python/distribute:tpu_strategy",
+        "//tensorflow/python/keras",
+        "//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
+    ],
+)
+
 cuda_py_test(
     name = "image_preprocessing_test",
     size = "medium",
@@ -318,6 +391,20 @@ tf_py_test(
     ],
 )
 
+tpu_py_test(
+    name = "normalization_tpu_test",
+    srcs = ["normalization_tpu_test.py"],
+    disable_experimental = True,
+    python_version = "PY3",
+    tags = ["no_oss"],
+    deps = [
+        ":normalization",
+        "//tensorflow/python/distribute:tpu_strategy",
+        "//tensorflow/python/keras",
+        "//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
+    ],
+)
+
 tf_py_test(
     name = "text_vectorization_test",
     size = "medium",
@@ -334,6 +421,20 @@ tf_py_test(
     ],
 )
 
+tpu_py_test(
+    name = "text_vectorization_tpu_test",
+    srcs = ["text_vectorization_tpu_test.py"],
+    disable_experimental = True,
+    python_version = "PY3",
+    tags = ["no_oss"],
+    deps = [
+        ":text_vectorization",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/keras",
+        "//tensorflow/python/keras/distribute:tpu_strategy_test_utils",
+    ],
+)
+
 tf_py_test(
     name = "reduction_test",
     srcs = ["reduction_test.py"],
@@ -345,22 +446,6 @@ tf_py_test(
     ],
 )
 
-tf_py_test(
-    name = "categorical_encoding_test",
-    size = "medium",
-    srcs = ["categorical_encoding_test.py"],
-    python_version = "PY3",
-    deps = [
-        ":categorical_encoding",
-        ":preprocessing_test_utils",
-        "//tensorflow/python:client_testlib",
-        "//tensorflow/python/keras",
-        "//tensorflow/python/keras/utils:generic_utils",
-        "//tensorflow/python/ops/ragged:ragged_string_ops",
-        "@absl_py//absl/testing:parameterized",
-    ],
-)
-
 tf_py_test(
     name = "preprocessing_stage_test",
     srcs = ["preprocessing_stage_test.py"],

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

@@ -0,0 +1,60 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for keras.layers.preprocessing.normalization."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.python import keras
+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)
+class CategoricalEncodingDistributionTest(
+    keras_parameterized.TestCase,
+    preprocessing_test_utils.PreprocessingLayerTest):
+
+  def test_tpu_distribution(self):
+    input_array = np.array([[1, 2, 3, 1], [0, 3, 1, 0]])
+
+    # pyformat: disable
+    expected_output = [[0, 1, 1, 1, 0, 0],
+                       [1, 1, 0, 1, 0, 0]]
+    # pyformat: enable
+    max_tokens = 6
+
+    strategy = tpu_strategy_test_utils.get_tpu_strategy()
+
+    with strategy.scope():
+      input_data = keras.Input(shape=(4,), dtype=dtypes.int32)
+      layer = categorical_encoding.CategoricalEncoding(
+          max_tokens=max_tokens, output_mode=categorical_encoding.BINARY)
+      int_data = layer(input_data)
+      model = keras.Model(inputs=input_data, outputs=int_data)
+    output_dataset = model.predict(input_array)
+    self.assertAllEqual(expected_output, output_dataset)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,57 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for keras.layers.preprocessing.normalization."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.python import keras
+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)
+class DiscretizationDistributionTest(
+    keras_parameterized.TestCase,
+    preprocessing_test_utils.PreprocessingLayerTest):
+
+  def test_tpu_distribution(self):
+    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():
+      input_data = keras.Input(shape=(None,))
+      layer = discretization.Discretization(
+          bins=[0., 1., 2.], output_mode=discretization.INTEGER)
+      bucket_data = layer(input_data)
+      self.assertAllEqual(expected_output_shape, bucket_data.shape.as_list())
+
+      model = keras.Model(inputs=input_data, outputs=bucket_data)
+    output_dataset = model.predict(input_array)
+    self.assertAllEqual(expected_output, output_dataset)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,58 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for keras.layers.preprocessing.normalization."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.python import keras
+from tensorflow.python.data.ops import dataset_ops
+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)
+class HashingDistributionTest(keras_parameterized.TestCase,
+                              preprocessing_test_utils.PreprocessingLayerTest):
+
+  def test_tpu_distribution(self):
+    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():
+      input_data = keras.Input(shape=(None,), dtype=dtypes.string)
+      layer = hashing.Hashing(num_bins=2)
+      int_data = layer(input_data)
+      model = keras.Model(inputs=input_data, outputs=int_data)
+    output_dataset = model.predict(input_dataset)
+    self.assertAllEqual(expected_output, output_dataset)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -30,6 +30,7 @@ from tensorflow.python import keras
 from tensorflow.python import tf2
 
 from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import one_device_strategy
 from tensorflow.python.eager import context
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import sparse_tensor
@@ -403,6 +404,29 @@ class CategoricalEncodingAdaptTest(
     layer.adapt(batched_ds)
 
 
+@keras_parameterized.run_all_keras_modes
+class IndexLookupDistributionTest(
+    keras_parameterized.TestCase,
+    preprocessing_test_utils.PreprocessingLayerTest):
+
+  def test_cpu_distribution(self):
+    vocab_data = ["earth", "wind", "and", "fire"]
+    input_array = np.array([["earth", "wind", "and", "fire"],
+                            ["fire", "and", "earth", "michigan"]])
+    expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
+
+    strategy = one_device_strategy.OneDeviceStrategy("/cpu:0")
+
+    with strategy.scope():
+      input_data = keras.Input(shape=(None,), dtype=dtypes.string)
+      layer = get_layer_class()()
+      layer.set_vocabulary(vocab_data)
+      int_data = layer(input_data)
+      model = keras.Model(inputs=input_data, outputs=int_data)
+    output_dataset = model.predict(input_array)
+    self.assertAllEqual(expected_output, output_dataset)
+
+
 @keras_parameterized.run_all_keras_modes
 class IndexLookupOutputTest(keras_parameterized.TestCase,
                             preprocessing_test_utils.PreprocessingLayerTest):

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

@@ -0,0 +1,66 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for keras.layers.preprocessing.normalization."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.python import keras
+from tensorflow.python.data.ops import dataset_ops
+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 preprocessing_test_utils
+from tensorflow.python.platform import test
+
+
+@keras_parameterized.run_all_keras_modes(
+    always_skip_v1=True, always_skip_eager=True)
+class IndexLookupDistributionTest(
+    keras_parameterized.TestCase,
+    preprocessing_test_utils.PreprocessingLayerTest):
+
+  def test_tpu_distribution(self):
+    vocab_data = [[
+        "earth", "earth", "earth", "earth", "wind", "wind", "wind", "and",
+        "and", "fire"
+    ]]
+    vocab_dataset = dataset_ops.Dataset.from_tensors(vocab_data)
+    input_array = np.array([["earth", "wind", "and", "fire"],
+                            ["fire", "and", "earth", "michigan"]])
+    input_dataset = dataset_ops.Dataset.from_tensor_slices(input_array).batch(
+        2, drop_remainder=True)
+    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():
+      input_data = keras.Input(shape=(None,), dtype=dtypes.string)
+      layer = index_lookup.IndexLookup()
+      layer.adapt(vocab_dataset)
+      int_data = layer(input_data)
+      model = keras.Model(inputs=input_data, outputs=int_data)
+    output_dataset = model.predict(input_dataset)
+    self.assertAllEqual(expected_output, output_dataset)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,128 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for keras.layers.preprocessing.normalization."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+
+import numpy as np
+
+from tensorflow.python import keras
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.keras import keras_parameterized
+from tensorflow.python.keras.distribute import tpu_strategy_test_utils
+from tensorflow.python.keras.layers.preprocessing import normalization
+from tensorflow.python.keras.layers.preprocessing import preprocessing_test_utils
+from tensorflow.python.platform import test
+
+
+def _get_layer_computation_test_cases():
+  test_cases = ({
+      "adapt_data": np.array([[1.], [2.], [3.], [4.], [5.]], dtype=np.float32),
+      "axis": -1,
+      "test_data": np.array([[1.], [2.], [3.]], np.float32),
+      "expected": np.array([[-1.414214], [-.707107], [0]], np.float32),
+      "testcase_name": "2d_single_element"
+  }, {
+      "adapt_data": np.array([[1.], [2.], [3.], [4.], [5.]], dtype=np.float32),
+      "axis": None,
+      "test_data": np.array([[1.], [2.], [3.]], np.float32),
+      "expected": np.array([[-1.414214], [-.707107], [0]], np.float32),
+      "testcase_name": "2d_single_element_none_axis"
+  }, {
+      "adapt_data": np.array([[1., 2., 3., 4., 5.]], dtype=np.float32),
+      "axis": None,
+      "test_data": np.array([[1.], [2.], [3.]], np.float32),
+      "expected": np.array([[-1.414214], [-.707107], [0]], np.float32),
+      "testcase_name": "2d_single_element_none_axis_flat_data"
+  }, {
+      "adapt_data":
+          np.array([[[1., 2., 3.], [2., 3., 4.]], [[3., 4., 5.], [4., 5., 6.]]],
+                   np.float32),
+      "axis":
+          1,
+      "test_data":
+          np.array([[[1., 2., 3.], [2., 3., 4.]], [[3., 4., 5.], [4., 5., 6.]]],
+                   np.float32),
+      "expected":
+          np.array([[[-1.549193, -0.774597, 0.], [-1.549193, -0.774597, 0.]],
+                    [[0., 0.774597, 1.549193], [0., 0.774597, 1.549193]]],
+                   np.float32),
+      "testcase_name":
+          "3d_internal_axis"
+  }, {
+      "adapt_data":
+          np.array(
+              [[[1., 0., 3.], [2., 3., 4.]], [[3., -1., 5.], [4., 5., 8.]]],
+              np.float32),
+      "axis": (1, 2),
+      "test_data":
+          np.array(
+              [[[3., 1., -1.], [2., 5., 4.]], [[3., 0., 5.], [2., 5., 8.]]],
+              np.float32),
+      "expected":
+          np.array(
+              [[[1., 3., -5.], [-1., 1., -1.]], [[1., 1., 1.], [-1., 1., 1.]]],
+              np.float32),
+      "testcase_name":
+          "3d_multiple_axis"
+  })
+
+  crossed_test_cases = []
+  # Cross above test cases with use_dataset in (True, False)
+  for use_dataset in (True, False):
+    for case in test_cases:
+      case = case.copy()
+      if use_dataset:
+        case["testcase_name"] = case["testcase_name"] + "_with_dataset"
+      case["use_dataset"] = use_dataset
+      crossed_test_cases.append(case)
+
+  return crossed_test_cases
+
+
+@keras_parameterized.run_all_keras_modes(
+    always_skip_v1=True, always_skip_eager=True)
+class NormalizationTest(keras_parameterized.TestCase,
+                        preprocessing_test_utils.PreprocessingLayerTest):
+
+  @parameterized.named_parameters(*_get_layer_computation_test_cases())
+  def test_layer_computation(self, adapt_data, axis, test_data, use_dataset,
+                             expected):
+    input_shape = tuple([None for _ in range(test_data.ndim - 1)])
+    if use_dataset:
+      # Keras APIs expect batched datasets
+      adapt_data = dataset_ops.Dataset.from_tensor_slices(adapt_data).batch(
+          test_data.shape[0] // 2)
+      test_data = dataset_ops.Dataset.from_tensor_slices(test_data).batch(
+          test_data.shape[0] // 2)
+
+    strategy = tpu_strategy_test_utils.get_tpu_strategy()
+
+    with strategy.scope():
+      input_data = keras.Input(shape=input_shape)
+      layer = normalization.Normalization(axis=axis)
+      layer.adapt(adapt_data)
+      output = layer(input_data)
+      model = keras.Model(input_data, output)
+      output_data = model.predict(test_data)
+    self.assertAllClose(expected, output_data)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -27,6 +27,7 @@ from tensorflow.python import keras
 from tensorflow.python import tf2
 
 from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import one_device_strategy
 from tensorflow.python.eager import context
 from tensorflow.python.framework import dtypes
 from tensorflow.python.keras import backend
@@ -568,6 +569,33 @@ class TextVectorizationPreprocessingTest(
     self.assertAllEqual(expected_output, output)
 
 
+@keras_parameterized.run_all_keras_modes
+class TextVectorizationDistributionTest(
+    keras_parameterized.TestCase,
+    preprocessing_test_utils.PreprocessingLayerTest):
+
+  def test_distribution_strategy_output(self):
+    vocab_data = ["earth", "wind", "and", "fire"]
+    input_array = np.array([["earth", "wind", "and", "fire"],
+                            ["fire", "and", "earth", "michigan"]])
+    expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
+
+    strategy = one_device_strategy.OneDeviceStrategy("/cpu:0")
+    with strategy.scope():
+      input_data = keras.Input(shape=(None,), dtype=dtypes.string)
+      layer = get_layer_class()(
+          max_tokens=None,
+          standardize=None,
+          split=None,
+          output_mode=text_vectorization.INT)
+      layer.set_vocabulary(vocab_data)
+      int_data = layer(input_data)
+      model = keras.Model(inputs=input_data, outputs=int_data)
+
+    output_dataset = model.predict(input_array)
+    self.assertAllEqual(expected_output, output_dataset)
+
+
 @keras_parameterized.run_all_keras_modes
 class TextVectorizationOutputTest(
     keras_parameterized.TestCase,

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

@@ -0,0 +1,96 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for keras.layers.preprocessing.normalization."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.python import keras
+from tensorflow.python.data.ops import dataset_ops
+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 preprocessing_test_utils
+from tensorflow.python.keras.layers.preprocessing import text_vectorization
+from tensorflow.python.platform import test
+
+
+@keras_parameterized.run_all_keras_modes(
+    always_skip_v1=True, always_skip_eager=True)
+class TextVectorizationTPUDistributionTest(
+    keras_parameterized.TestCase,
+    preprocessing_test_utils.PreprocessingLayerTest):
+
+  def test_distribution_strategy_output(self):
+    vocab_data = ["earth", "wind", "and", "fire"]
+    input_array = np.array([["earth", "wind", "and", "fire"],
+                            ["fire", "and", "earth", "michigan"]])
+    input_dataset = dataset_ops.Dataset.from_tensor_slices(input_array).batch(
+        2, drop_remainder=True)
+
+    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():
+      input_data = keras.Input(shape=(None,), dtype=dtypes.string)
+      layer = text_vectorization.TextVectorization(
+          max_tokens=None,
+          standardize=None,
+          split=None,
+          output_mode=text_vectorization.INT)
+      layer.set_vocabulary(vocab_data)
+      int_data = layer(input_data)
+      model = keras.Model(inputs=input_data, outputs=int_data)
+
+    output_dataset = model.predict(input_dataset)
+    self.assertAllEqual(expected_output, output_dataset)
+
+  def test_distribution_strategy_output_with_adapt(self):
+    vocab_data = [[
+        "earth", "earth", "earth", "earth", "wind", "wind", "wind", "and",
+        "and", "fire"
+    ]]
+    vocab_dataset = dataset_ops.Dataset.from_tensors(vocab_data)
+    input_array = np.array([["earth", "wind", "and", "fire"],
+                            ["fire", "and", "earth", "michigan"]])
+    input_dataset = dataset_ops.Dataset.from_tensor_slices(input_array).batch(
+        2, drop_remainder=True)
+
+    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():
+      input_data = keras.Input(shape=(None,), dtype=dtypes.string)
+      layer = text_vectorization.TextVectorization(
+          max_tokens=None,
+          standardize=None,
+          split=None,
+          output_mode=text_vectorization.INT)
+      layer.adapt(vocab_dataset)
+      int_data = layer(input_data)
+      model = keras.Model(inputs=input_data, outputs=int_data)
+
+    output_dataset = model.predict(input_dataset)
+    self.assertAllEqual(expected_output, output_dataset)
+if __name__ == "__main__":
+  test.main()