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:
parent
232fb86bfb
commit
c9e11b03d1
@ -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()
|
||||
|
Loading…
Reference in New Issue
Block a user