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

@@ -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"],