Switch custom ops (#2397)

* Add methods for switching custom kernels

* Motivate usage in codebase

* -custom_kernel_disabled

* - doc fix

* fix readme
fix code doc resource_loader.py
fix code doc options.py

* is_custom_disabled

* - is_custom_kernel_disabled()

* - rm disable from gelu.py

* - rm pass from options.py

* - change docs of options.py disable & enable

* - options.py  is_custom_kernel docfix

* - doc fix resource_loader.py

* - fix README.md doc

* - previous_py_ops_value test_utils.py

* - format README.md

* - rm TF_ADDONS_PY_OPS in distort_image_ops.py

* Update README.md

Co-authored-by: Tzu-Wei Sung <[email protected]>

* Update README.md

Co-authored-by: Tzu-Wei Sung <[email protected]>

* Update tensorflow_addons/options.py

Co-authored-by: Tzu-Wei Sung <[email protected]>

* Update tensorflow_addons/options.py

Co-authored-by: Tzu-Wei Sung <[email protected]>

* Update tensorflow_addons/utils/resource_loader.py

Co-authored-by: Tzu-Wei Sung <[email protected]>

* Update tensorflow_addons/utils/test_utils.py

Co-authored-by: Tzu-Wei Sung <[email protected]>

* Update tensorflow_addons/options.py

Co-authored-by: Tzu-Wei Sung <[email protected]>

* Fix typo and format

* Expose options to public

* Hide some internal uses

* Wording

* Wording

Co-authored-by: Tzu-Wei Sung <[email protected]>

Author: jonpsy

README.md

@@ -194,9 +194,9 @@ The order of priority on Linux is:
 2) C++ implementation
 3) Pure TensorFlow + Python implementation (works on CPU and GPU)
 
-If you want to change the default priority, "C++ and CUDA" VS "pure TensorFlow Python", 
-you can set the variable `TF_ADDONS_PY_OPS` either from the command line or in 
-your code.
+If you want to change the default priority, "C++ and CUDA" VS "pure TensorFlow Python",
+you can set the environment variable `TF_ADDONS_PY_OPS=1` from the command line or
+run `tfa.options.disable_custom_kernel()` in your code.
 
 For example, if you are on Linux and you have compatibility problems with the compiled ops,
 you can give priority to the Python implementations:
@@ -210,7 +210,7 @@ or in your code:
 
 ```python
 import tensorflow_addons as tfa
-tfa.options.TF_ADDONS_PY_OPS = True
+tfa.options.disable_custom_kernel()
 ```
 
 This variable defaults to `True` on Windows and macOS, and `False` on Linux.

tensorflow_addons/__init__.py

@@ -28,6 +28,7 @@
 from tensorflow_addons import rnn
 from tensorflow_addons import seq2seq
 from tensorflow_addons import text
+from tensorflow_addons import options
 from tensorflow_addons.register import register_all
 from tensorflow_addons.utils import types
 

tensorflow_addons/image/distort_image_ops.py

@@ -183,7 +183,7 @@ def adjust_hsv_in_yiq(
         )
         scale_value = tf.cast(scale_value, dtype=image.dtype, name="scale_value")
 
-        if not options.TF_ADDONS_PY_OPS:
+        if not options.is_custom_kernel_disabled():
             warnings.warn(
                 "C++/CUDA kernel of `adjust_hsv_in_yiq` will be removed in Addons `0.13`.",
                 DeprecationWarning,

tensorflow_addons/options.py

@@ -4,15 +4,14 @@
 import traceback
 
 try:
-    TF_ADDONS_PY_OPS = bool(int(os.environ["TF_ADDONS_PY_OPS"]))
+    _TF_ADDONS_PY_OPS = bool(int(os.environ["TF_ADDONS_PY_OPS"]))
 except KeyError:
     if platform.system() == "Linux":
-        TF_ADDONS_PY_OPS = False
+        _TF_ADDONS_PY_OPS = False
     else:
-        TF_ADDONS_PY_OPS = True
+        _TF_ADDONS_PY_OPS = True
 
-
-FALLBACK_WARNING_TEMPLATE = """{}
+_FALLBACK_WARNING_TEMPLATE = """{}
 
 The {} C++/CUDA custom op could not be loaded.
 For this reason, Addons will fallback to an implementation written
@@ -26,24 +25,47 @@
 If you want this warning to disappear, either make sure the TensorFlow installed
 is compatible with this version of Addons, or tell TensorFlow Addons to
 prefer using Python implementations and not custom C++/CUDA ones. You can do that
-by changing the TF_ADDONS_PY_OPS flag
-either with the environment variable:
+by setting the enviornment variable `TF_ADDONS_PY_OPS=1`:
 ```bash
 TF_ADDONS_PY_OPS=1 python my_script.py
 ```
-or in your code, after your imports:
+or run `tfa.options.disable_custom_kernel()` in your code, after your imports:
 ```python
 import tensorflow_addons as tfa
 import ...
 import ...
 
-tfa.options.TF_ADDONS_PY_OPS = True
+tfa.options.disable_custom_kernel()
 ```
 """
 
 
 def warn_fallback(op_name):
-    warning_msg = FALLBACK_WARNING_TEMPLATE.format(traceback.format_exc(), op_name)
+    warning_msg = _FALLBACK_WARNING_TEMPLATE.format(traceback.format_exc(), op_name)
     warnings.warn(warning_msg, RuntimeWarning)
-    global TF_ADDONS_PY_OPS
-    TF_ADDONS_PY_OPS = True
+    disable_custom_kernel()
+
+
+def enable_custom_kernel():
+    """Prefer custom C++/CUDA kernel to pure python operations.
+
+    Enable using custom C++/CUDA kernel instead of pure python operations.
+    It has the same effect as setting environment variable `TF_ADDONS_PY_OPS=0`.
+    """
+    global _TF_ADDONS_PY_OPS
+    _TF_ADDONS_PY_OPS = False
+
+
+def disable_custom_kernel():
+    """Prefer pure python operations to custom C++/CUDA kernel.
+
+    Disable using custom C++/CUDA kernel instead of pure python operations.
+    It has the same effect as setting environment variable `TF_ADDONS_PY_OPS=1`.
+    """
+    global _TF_ADDONS_PY_OPS
+    _TF_ADDONS_PY_OPS = True
+
+
+def is_custom_kernel_disabled():
+    """Return whether custom C++/CUDA kernel is disabled."""
+    return _TF_ADDONS_PY_OPS

tensorflow_addons/seq2seq/beam_search_decoder.py

@@ -234,7 +234,7 @@ def gather_tree(
     Raises:
       InvalidArgumentError: if `parent_ids` contains an invalid index.
     """
-    if not options.TF_ADDONS_PY_OPS:
+    if not options.is_custom_kernel_disabled():
         try:
             return _beam_search_so.ops.addons_gather_tree(
                 step_ids, parent_ids, max_sequence_lengths, end_token

tensorflow_addons/seq2seq/tests/beam_search_ops_test.py

@@ -83,7 +83,7 @@ def test_bad_parent_values_on_gpu():
     expected_result = _transpose_batch_time(
         [[[2, -1, 2], [6, 5, 6], [7, 8, 9], [10, 10, 10]]]
     )
-    if options.TF_ADDONS_PY_OPS:
+    if options.is_custom_kernel_disabled():
         # The Python version has the same behavior on CPU and GPU.
         with pytest.raises(tf.errors.InvalidArgumentError, match="parent id"):
             _ = gather_tree(

tensorflow_addons/utils/resource_loader.py

@@ -86,7 +86,8 @@ def display_warning_if_incompatible(self):
             "on Github. This is a known limitation."
             "\n\n"
             "It might help you to fallback to pure Python "
-            "ops with TF_ADDONS_PY_OPS . To do that, see "
+            "ops by setting environment variable `TF_ADDONS_PY_OPS=1` or using `tfa.options.disable_custom_kernel()` in your code. "
+            "To do that, see "
             "https://github.com/tensorflow/addons#gpucpu-custom-ops "
             "\n\n"
             "You can also change the TensorFlow version installed on your system. "

tensorflow_addons/utils/test_utils.py

@@ -95,14 +95,17 @@ def only_run_functions_eagerly(request):
 
 @pytest.fixture(scope="function", params=["custom_ops", "py_ops"])
 def run_custom_and_py_ops(request):
-    previous_py_ops_value = options.TF_ADDONS_PY_OPS
+    previous_is_custom_kernel_disabled = options.is_custom_kernel_disabled()
     if request.param == "custom_ops":
-        options.TF_ADDONS_PY_OPS = False
+        options.enable_custom_kernel()
     elif request.param == "py_ops":
-        options.TF_ADDONS_PY_OPS = True
+        options.disable_custom_kernel()
 
     def _restore_py_ops_value():
-        options.TF_ADDONS_PY_OPS = previous_py_ops_value
+        if previous_is_custom_kernel_disabled:
+            options.disable_custom_kernel()
+        else:
+            options.enable_custom_kernel()
 
     request.addfinalizer(_restore_py_ops_value)
 

tools/docs/build_docs.py

@@ -73,7 +73,10 @@ def main(argv):
         root_title=PROJECT_FULL_NAME,
         py_modules=[(PROJECT_SHORT_NAME, tfa)],
         code_url_prefix=code_url_prefix,
-        private_map={"tfa": ["__version__", "utils", "version"]},
+        private_map={
+            "tfa": ["__version__", "utils", "version"],
+            "tfa.options": ["warn_fallback"],
+        },
         # These callbacks usually clean up a lot of aliases caused by internal imports.
         callbacks=[
             public_api.local_definitions_filter,