Repository: tensorflow/addons

Issue:
Public API to force load custom ops
Currently, it is inconvenient to load in Python a SavedModel that includes Addons custom ops. Consider the example below:

* `save.py`

```python
import tensorflow as tf
import tensorflow_addons as tfa

class Model(tf.keras.Model):
    @tf.function(input_signature=(tf.TensorSpec(shape=[None, 32], dtype=tf.float32),))
    def call(self, x):
        return tfa.activations.gelu(x)

model = Model()
tf.saved_model.save(model, '/tmp/model', signatures=model.call)
```

* `load.py`

```python
import tensorflow as tf
tf.saved_model.load("/tmp/model")
```

The load will fail because Addons custom ops are not registered to the TensorFlow runtime. This is expected as we first have to invoke `tf.load_op_library` on the custom ops.

However, with the new work on lazy loading #855 it got harder to force this op registration. For this model, the user should run the following which relies on internal APIs:

```python
from tensorflow_addons.activations.gelu import _activation_so
_activation_so.ops
```

If the custom ops are not loaded during the main import (i.e. during `import tensorflow_addons`), then the package should expose a public API that registers all custom ops.

Any thoughts?


Discussion:
That's definitely a big problem.

I'd be in favor of avoiding using an explicit function to force the loading of the ops. It's not a very good move for UX as users might expect this to be automatic.

We can try to load all the SO at import time and then throw a warning if something goes wrong? With a mechanism similar to #1137 . But then we should have a way of disabling the warnings if some users don't care about custom ops.

In a perfect world, one would be able to save a model using the .so, and another user could reload the same model using the python-only equivalent op. Is that possible? Maybe with the keras model saving?

EDIT: Actually, I'm not so sure it's the best UX to register everything at import time. If we made a function to register all ops, would we make this function register keras functions too? The `register_keras_serializable` part?
From the user perspective, loading everything at import time sounds like the less surprising behavior. In the example above,

```python
import tensorflow as tf
tf.saved_model.load("/tmp/model")
```

adding the Addons import would seem a natural fix for the SavedModel loading issue:

```python
import tensorflow as tf
import tensorflow_addons as tfa
tf.saved_model.load("/tmp/model")
```

> We can try to load all the SO at import time and then throw a warning if something goes wrong?

IIRC, the custom op errors we faced was mostly segmentation faults on import, right? In that case, it would be difficult to catch the error. We could fork the process and load the custom op but that seems overkill.

Or we could load custom ops by default and add an option to disable this automatic loading for users facing issues.
When loading a .so, if there are any kind of issue, it throws a `tensorflow.errors.NotFoundError`, not a segfault. So it's easy to give a warning and just continue if there is a problem at import time.

I believe we need to discuss more this statement: 
> From the user perspective, loading everything at import time sounds like the less surprising behavior. 

If we take the "load the .so files at import time" approach:

```python
import tensorflow as tf
import tensorflow_addons as tfa
tf.saved_model.load("/tmp/model")
```

If the user is using pycharm, it's going to tell the user: `tensorflow_addons imported but unused`
If the user is using visual studio code, it's going to tell the user: `tensorflow_addons imported but unused`
If the user is using flake8, it's going to tell the user: `tensorflow_addons imported but unused`
If the user has a colleague who has never worked with addons, it's going to say that tensorflow_addons was imported but unused and add a commit to remove the import.

This doesn't happen if we go for:

```python
import tensorflow as tf
import tensorflow_addons as tfa
tfa.register_all_ops()
tf.saved_model.load("/tmp/model")
```

When reading the code done with a library that has a well written API (like keras/pygithub/numpy...)  it's possible to follow a piece of code without reading the documentation beforehand. In this case it's really clear what is happening, unlike when the shared objects are loaded at import time. A piece of code is read many more time than it's written, so we should go for readability.

If we follow the principal of least surprises, for some users, doing any significant work at import time  sounds very strange and unintuitive, as demonstrated by the errors thrown by flake8, pycharm and visual studio code. It's also against the zen of python:

> Explicit is better than implicit

We have other benefits in doing this: we catch errors early and we don't make an additional environment variable for people who don't want to load the SO files or want to debug what happens at import time. Because the function call means "load everything" and not "maybe load everything" we can throw a hard error and it's easier to debug. People who want to use TFA with another version of Tensorflow than the recommended one can just not use the function and they won't have issues with the shared objects.

I would recommend making a new function as proposed by @guillaumekln in his first message:

```python
def register_all_ops(keras_objects=True, custom_kernels=True):
    if keras_objects:
        for obj in keras_objects:
            tf.keras.register_keras_serializable(obj)
    if custom_kernels:
        for custom_kernel_path in custom_kernel_paths:
            tf.load_op_library(custom_kernel_path)
```



Resolving patch:
diff --git a/tensorflow_addons/__init__.py b/tensorflow_addons/__init__.py
--- a/tensorflow_addons/__init__.py
+++ b/tensorflow_addons/__init__.py
@@ -28,5 +28,6 @@
 from tensorflow_addons import rnn
 from tensorflow_addons import seq2seq
 from tensorflow_addons import text
+from tensorflow_addons.register import register_all
 
 from tensorflow_addons.version import __version__
diff --git a/tensorflow_addons/register.py b/tensorflow_addons/register.py
new file mode 100644
--- /dev/null
+++ b/tensorflow_addons/register.py
@@ -0,0 +1,109 @@
+import glob
+import os
+from pathlib import Path
+
+import tensorflow as tf
+
+from tensorflow_addons.utils.resource_loader import get_project_root
+
+
+def register_all(keras_objects: bool = True, custom_kernels: bool = True) -> None:
+    """Register TensorFlow Addons' objects in TensorFlow global dictionaries.
+
+    When loading a Keras model that has a TF Addons' function, it is needed
+    for this function to be known by the Keras deserialization process.
+
+    There are two ways to do this, either do
+
+    ```python
+    tf.keras.models.load_model(
+        "my_model.tf",
+        custom_objects={"LAMB": tfa.image.optimizer.LAMB}
+    )
+    ```
+
+    or you can do:
+    ```python
+    tfa.register_all()
+    tf.tf.keras.models.load_model("my_model.tf")
+    ```
+
+    If the model contains custom ops (compiled ops) of TensorFlow Addons,
+    and the graph is loaded with `tf.saved_model.load`, then custom ops need
+    to be registered before to avoid an error of the type:
+
+    ```
+    tensorflow.python.framework.errors_impl.NotFoundError: Op type not registered
+    '...' in binary running on ... Make sure the Op and Kernel are
+    registered in the binary running in this process.
+    ```
+
+    In this case, the only way to make sure that the ops are registered is to call
+    this function:
+
+    ```python
+    tfa.register_all()
+    tf.saved_model.load("my_model.tf")
+    ```
+
+    Note that you can call this function multiple times in the same process,
+    it only has an effect the first time. Afterward, it's just a no-op.
+
+    Args:
+        keras_objects: boolean, `True` by default. If `True`, register all
+            Keras objects
+            with `tf.keras.utils.register_keras_serializable(package="Addons")`
+            If set to False, doesn't register any Keras objects
+            of Addons in TensorFlow.
+        custom_kernels: boolean, `True` by default. If `True`, loads all
+            custom kernels of TensorFlow Addons with
+            `tf.load_op_library("path/to/so/file.so")`. Loading the SO files
+            register them automatically. If `False` doesn't load and register
+            the shared objects files. Not that it might be useful to turn it off
+            if your installation of Addons doesn't work well with custom ops.
+    Returns:
+        None
+    """
+    if keras_objects:
+        register_keras_objects()
+    if custom_kernels:
+        register_custom_kernels()
+
+
+def register_keras_objects() -> None:
+    # TODO: once layer_test is replaced by a public API
+    # and we can used unregistered objects with it
+    # we can remove all decorators.
+    # And register Keras objects here.
+    pass
+
+
+def register_custom_kernels() -> None:
+    all_shared_objects = _get_all_shared_objects()
+    if not all_shared_objects:
+        raise FileNotFoundError(
+            "No shared objects files were found in the custom ops "
+            "directory in Tensorflow Addons, check your installation again,"
+            "or, if you don't need custom ops, call `tfa.register_all(custom_kernels=False)`"
+            " instead."
+        )
+    try:
+        for shared_object in all_shared_objects:
+            tf.load_op_library(shared_object)
+    except tf.errors.NotFoundError as e:
+        raise RuntimeError(
+            "One of the shared objects ({}) could not be loaded. This may be "
+            "due to a number of reasons (incompatible TensorFlow version, buiding from "
+            "source with different flags, broken install of TensorFlow Addons...). If you"
+            "wanted to register the shared objects because you needed them when loading your "
+            "model, you should fix your install of TensorFlow Addons. If you don't "
+            "use custom ops in your model, you can skip registering custom ops with "
+            "`tfa.register_all(custom_kernels=False)`".format(shared_object)
+        ) from e
+
+
+def _get_all_shared_objects():
+    custom_ops_dir = os.path.join(get_project_root(), "custom_ops")
+    all_shared_objects = glob.glob(custom_ops_dir + "/**/*.so", recursive=True)
+    all_shared_objects = [x for x in all_shared_objects if Path(x).is_file()]
+    return all_shared_objects
diff --git a/tools/ci_build/verify/check_typing_info.py b/tools/ci_build/verify/check_typing_info.py
--- a/tools/ci_build/verify/check_typing_info.py
+++ b/tools/ci_build/verify/check_typing_info.py
@@ -13,12 +13,9 @@
 # limitations under the License.
 # ==============================================================================
 #
-
-from types import ModuleType
-
 from typedapi import ensure_api_is_typed
 
-import tensorflow_addons
+import tensorflow_addons as tfa
 
 TUTORIAL_URL = "https://docs.python.org/3/library/typing.html"
 HELP_MESSAGE = (
@@ -30,11 +27,18 @@
 EXCEPTION_LIST = []
 
 
-modules_list = []
-for attr_name in dir(tensorflow_addons):
-    attr = getattr(tensorflow_addons, attr_name)
-    if isinstance(attr, ModuleType) and attr is not tensorflow_addons.options:
-        modules_list.append(attr)
+modules_list = [
+    tfa,
+    tfa.activations,
+    tfa.callbacks,
+    tfa.image,
+    tfa.losses,
+    tfa.metrics,
+    tfa.optimizers,
+    tfa.rnn,
+    tfa.seq2seq,
+    tfa.text,
+]
 
 
 if __name__ == "__main__":


Test patch:
diff --git a/tensorflow_addons/register_test.py b/tensorflow_addons/register_test.py
new file mode 100644
--- /dev/null
+++ b/tensorflow_addons/register_test.py
@@ -0,0 +1,20 @@
+import unittest
+import tensorflow as tf
+from tensorflow_addons.register import register_all, _get_all_shared_objects
+
+
+class AssertRNNCellTest(unittest.TestCase):
+    def test_multiple_register(self):
+        register_all()
+        register_all()
+
+    def test_get_all_shared_objects(self):
+        all_shared_objects = _get_all_shared_objects()
+        self.assertTrue(len(all_shared_objects) >= 4)
+
+        for file in all_shared_objects:
+            tf.load_op_library(file)
+
+
+if __name__ == "__main__":
+    unittest.main()