From 8dd818b0f60d0cc2cbd432994d3d82f84f2982b1 Mon Sep 17 00:00:00 2001
From: Chenkai Kuang <chenkai@google.com>
Date: Wed, 27 May 2020 14:28:30 -0700
Subject: [PATCH] Add gpu doc_test that covers all modules under
 tensorflow.python.distribute excepts tpu_strategy.

PiperOrigin-RevId: 313464425
Change-Id: I7f691e66f76d1404b1c5763eb85182e598c85703
---
 tensorflow/tools/docs/BUILD | 32 ++++++++++++++++++++++++++++++--
 1 file changed, 30 insertions(+), 2 deletions(-)

diff --git a/tensorflow/tools/docs/BUILD b/tensorflow/tools/docs/BUILD
index c0442a5986d..1f2dd5d31d2 100644
--- a/tensorflow/tools/docs/BUILD
+++ b/tensorflow/tools/docs/BUILD
@@ -11,6 +11,8 @@ package(
 
 exports_files(["LICENSE"])
 
+tpu_module = "tpu.,distribute.tpu_strategy"
+
 py_library(
     name = "tf_doctest_lib",
     srcs = ["tf_doctest_lib.py"],
@@ -23,7 +25,7 @@ py_library(
 py_test(
     name = "tf_doctest",
     srcs = ["tf_doctest.py"],
-    args = ["--module_prefix_skip=tpu.,distribute.tpu_strategy"],
+    args = ["--module_prefix_skip=" + tpu_module],
     python_version = "PY3",
     tags = [
         "no_oss_py2",
@@ -45,7 +47,7 @@ py_test(
 tpu_py_test(
     name = "tf_doctest_tpu",
     srcs = ["tf_doctest.py"],
-    args = ["--module=tpu.,distribute.tpu_strategy"],
+    args = ["--module=" + tpu_module],
     disable_experimental = True,
     disable_v3 = True,
     main = "tf_doctest.py",
@@ -64,6 +66,32 @@ tpu_py_test(
     ],
 )
 
+py_test(
+    name = "tf_doctest_gpu",
+    srcs = ["tf_doctest.py"],
+    args = [
+        "--module=distribute.",
+        "--module_prefix_skip=" + tpu_module,
+    ],
+    main = "tf_doctest.py",
+    python_version = "PY3",
+    tags = [
+        "no_oss_py2",
+        "no_pip",
+        "no_rocm",
+        "no_windows",  # numpy prints differently on windows.
+        "noasan",
+        "nomsan",
+        "notsan",
+        "requires-gpu-nvidia",
+    ],
+    deps = [
+        ":tf_doctest_lib",
+        "//tensorflow:tensorflow_py",
+        "//third_party/py/numpy",
+    ],
+)
+
 py_test(
     name = "tf_doctest_test",
     srcs = ["tf_doctest_test.py"],