Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow_addons/options.py: 67%

1import os 

2import platform 

3import warnings 

4import traceback 

5 

6try: 

7 _TF_ADDONS_PY_OPS = bool(int(os.environ["TF_ADDONS_PY_OPS"])) 

8except KeyError: 

9 if platform.system() == "Linux": 

10 _TF_ADDONS_PY_OPS = False 

11 else: 

12 _TF_ADDONS_PY_OPS = True 

13 

14_FALLBACK_WARNING_TEMPLATE = """{} 

15 

16The {} C++/CUDA custom op could not be loaded. 

17For this reason, Addons will fallback to an implementation written 

18in Python with public TensorFlow ops. There worst you might experience with 

19this is a moderate slowdown on GPU. There can be multiple 

20reason for this loading error, one of them may be an ABI incompatibility between 

21the TensorFlow installed on your system and the TensorFlow used to compile 

22TensorFlow Addons' custom ops. The stacktrace generated when loading the 

23shared object file was displayed above. 

24 

25If you want this warning to disappear, either make sure the TensorFlow installed 

26is compatible with this version of Addons, or tell TensorFlow Addons to 

27prefer using Python implementations and not custom C++/CUDA ones. You can do that 

28by setting the enviornment variable `TF_ADDONS_PY_OPS=1`: 

29```bash 

30TF_ADDONS_PY_OPS=1 python my_script.py 

31``` 

32or run `tfa.options.disable_custom_kernel()` in your code, after your imports: 

33```python 

34import tensorflow_addons as tfa 

35import ... 

36import ... 

37 

38tfa.options.disable_custom_kernel() 

39``` 

40""" 

41 

42 

43def warn_fallback(op_name): 

44 warning_msg = _FALLBACK_WARNING_TEMPLATE.format(traceback.format_exc(), op_name) 

45 warnings.warn(warning_msg, RuntimeWarning) 

46 disable_custom_kernel() 

47 

48 

49def enable_custom_kernel(): 

50 """Prefer custom C++/CUDA kernel to pure python operations. 

51 

52 Enable using custom C++/CUDA kernel instead of pure python operations. 

53 It has the same effect as setting environment variable `TF_ADDONS_PY_OPS=0`. 

54 """ 

55 global _TF_ADDONS_PY_OPS 

56 _TF_ADDONS_PY_OPS = False 

57 

58 

59def disable_custom_kernel(): 

60 """Prefer pure python operations to custom C++/CUDA kernel. 

61 

62 Disable using custom C++/CUDA kernel instead of pure python operations. 

63 It has the same effect as setting environment variable `TF_ADDONS_PY_OPS=1`. 

64 """ 

65 global _TF_ADDONS_PY_OPS 

66 _TF_ADDONS_PY_OPS = True 

67 

68 

69def is_custom_kernel_disabled(): 

70 """Return whether custom C++/CUDA kernel is disabled.""" 

71 return _TF_ADDONS_PY_OPS