Treat "test_xla_gpu" as GPU_TEST in "NamedGPUCombination".

PiperOrigin-RevId: 309276564
Change-Id: Iea914898c6a090c8397d0ad0890e239948eec943
This commit is contained in:
Chenkai Kuang 2020-04-30 12:25:37 -07:00 committed by TensorFlower Gardener
parent f69af82c3d
commit dfe0f0aa36
1 changed files with 3 additions and 2 deletions

View File

@ -22,6 +22,7 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import re
import sys import sys
import types import types
import unittest import unittest
@ -94,10 +95,10 @@ class NamedGPUCombination(combinations_lib.TestCombination):
Attributes: Attributes:
GPU_TEST: The environment is considered to have GPU hardware available if GPU_TEST: The environment is considered to have GPU hardware available if
the name of the program contains "test_gpu". the name of the program contains "test_gpu" or "test_xla_gpu".
""" """
GPU_TEST = "test_gpu" in sys.argv[0] GPU_TEST = re.search(r"(test_gpu|test_xla_gpu)$", sys.argv[0])
def should_execute_combination(self, kwargs): def should_execute_combination(self, kwargs):
distributions = [ distributions = [