From b646d41f69f8a1cf2f5c76d934adbf8c067e4444 Mon Sep 17 00:00:00 2001
From: Hye Soo Yang <hyey@google.com>
Date: Wed, 20 Nov 2019 07:49:33 -0800
Subject: [PATCH] Add `python_version` arg to python test macros in
 `tensorflow.bzl`.

PiperOrigin-RevId: 281518994
Change-Id: I14b8daf75cd75776a23864236f5b9f3262904385
---
 tensorflow/tensorflow.bzl | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl
index 57ebb1e2164..2376f0d3988 100644
--- a/tensorflow/tensorflow.bzl
+++ b/tensorflow/tensorflow.bzl
@@ -2152,7 +2152,8 @@ def gpu_py_test(
         flaky = 0,
         xla_enable_strict_auto_jit = False,
         xla_enabled = False,
-        grpc_enabled = False):
+        grpc_enabled = False,
+        python_version = "PY2"):
     # TODO(b/122522101): Don't ignore xla_enable_strict_auto_jit and enable additional
     # XLA tests once enough compute resources are available.
     _ignored = [xla_enable_strict_auto_jit]
@@ -2179,6 +2180,7 @@ def gpu_py_test(
             tags = test_tags,
             xla_enabled = xla_enabled,
             xla_enable_strict_auto_jit = False,
+            python_version = python_version,
         )
 
 register_extension_info(
@@ -2243,7 +2245,8 @@ def py_tests(
         prefix = "",
         xla_enable_strict_auto_jit = False,
         xla_enabled = False,
-        grpc_enabled = False):
+        grpc_enabled = False,
+        python_version = "PY2"):
     for src in srcs:
         test_name = src.split("/")[-1].split(".")[0]
         if prefix:
@@ -2261,6 +2264,7 @@ def py_tests(
             tags = tags,
             xla_enabled = xla_enabled,
             xla_enable_strict_auto_jit = xla_enable_strict_auto_jit,
+            python_version = python_version,
         )
 
 def gpu_py_tests(
@@ -2275,7 +2279,8 @@ def gpu_py_tests(
         prefix = "",
         xla_enable_strict_auto_jit = False,
         xla_enabled = False,
-        grpc_enabled = False):
+        grpc_enabled = False,
+        python_version = "PY2"):
     # TODO(b/122522101): Don't ignore xla_enable_strict_auto_jit and enable additional
     # XLA tests once enough compute resources are available.
     _ignored = [xla_enable_strict_auto_jit]
@@ -2293,6 +2298,7 @@ def gpu_py_tests(
         tags = test_tags,
         xla_enabled = xla_enabled,
         xla_enable_strict_auto_jit = False,
+        python_version = python_version,
     )
 
 # terminology changes: saving cuda_* definition for compatibility