This commit adds a test.is_built_with_rocm python utility routine, which will return true/false based on whether TF was built with ROCm support enables (i.e. with --config=rocm )

This routine will be used in subsequent PRs to make ROCm specific changes to certain unit tests. The nature of these changes will be to primarily disable certain subtests, which test features that are not yet supported on the ROCm platform (for e.g. complex datatypes, 3D pooling, etc)

A `test.is_built_with_gpu_support` routine is also being added, which is essentially "is_built_with_cuda() or is_built_with_rocm()".
This commit is contained in:
Deven Desai 2019-06-26 21:20:23 +00:00
parent 232fb86bfb
commit c9e11b03d1

View File

@ -94,3 +94,15 @@ def test_src_dir_path(relative_path):
def is_built_with_cuda():
"""Returns whether TensorFlow was built with CUDA (GPU) support."""
return _test_util.IsGoogleCudaEnabled()
@tf_export('test.is_built_with_rocm')
def is_built_with_rocm():
"""Returns whether TensorFlow was built with ROCm (GPU) support."""
return _test_util.IsBuiltWithROCm()
@tf_export('test.is_built_with_gpu_support')
def is_built_with_gpu_support():
"""Returns whether TensorFlow was built with GPU (either CUDA or ROCm) support.
"""
return is_built_with_cuda() or is_built_with_rocm()