[WIP] Use public Keras object registration (#669)

* Remove test skips since upstream fix

* Use public keras object registration

* Merge master and update registration

* Fix tf import

* Update READMEs

* F scores keras registration

* Lint

Author: seanpmorgan

tensorflow_addons/activations/README.md

@@ -29,8 +29,7 @@
 In order to conform with the current API standard, all activations
 must:
  * Be a `tf.function`.
- * [Register as a keras global object](https://github.com/tensorflow/addons/blob/master/tensorflow_addons/utils/keras_utils.py)
-  so it can be serialized properly.
+ * Register as a keras global object so it can be serialized properly: `@tf.keras.utils.register_keras_serializable(package='Addons')`
  * Add the addon to the `py_library` in this sub-package's BUILD file.
 
 #### Testing Requirements

tensorflow_addons/activations/gelu.py

@@ -18,14 +18,13 @@
 from __future__ import print_function
 
 import tensorflow as tf
-from tensorflow_addons.utils import keras_utils
 from tensorflow_addons.utils.resource_loader import get_path_to_datafile
 
 _activation_ops_so = tf.load_op_library(
     get_path_to_datafile("custom_ops/activations/_activation_ops.so"))
 
 
-@keras_utils.register_keras_custom_object
+@tf.keras.utils.register_keras_serializable(package='Addons')
 @tf.function
 def gelu(x, approximate=True):
     """Gaussian Error Linear Unit.

tensorflow_addons/activations/hardshrink.py

@@ -18,14 +18,13 @@
 from __future__ import print_function
 
 import tensorflow as tf
-from tensorflow_addons.utils import keras_utils
 from tensorflow_addons.utils.resource_loader import get_path_to_datafile
 
 _activation_ops_so = tf.load_op_library(
     get_path_to_datafile("custom_ops/activations/_activation_ops.so"))
 
 
-@keras_utils.register_keras_custom_object
+@tf.keras.utils.register_keras_serializable(package='Addons')
 @tf.function
 def hardshrink(x, lower=-0.5, upper=0.5):
     """Hard shrink function.

tensorflow_addons/activations/lisht.py

@@ -18,14 +18,13 @@
 from __future__ import print_function
 
 import tensorflow as tf
-from tensorflow_addons.utils import keras_utils
 from tensorflow_addons.utils.resource_loader import get_path_to_datafile
 
 _activation_ops_so = tf.load_op_library(
     get_path_to_datafile("custom_ops/activations/_activation_ops.so"))
 
 
-@keras_utils.register_keras_custom_object
+@tf.keras.utils.register_keras_serializable(package='Addons')
 @tf.function
 def lisht(x):
     """LiSHT: Non-Parameteric Linearly Scaled Hyperbolic Tangent Activation Function.

tensorflow_addons/activations/mish.py

@@ -18,14 +18,13 @@
 from __future__ import print_function
 
 import tensorflow as tf
-from tensorflow_addons.utils import keras_utils
 from tensorflow_addons.utils.resource_loader import get_path_to_datafile
 
 _activation_ops_so = tf.load_op_library(
     get_path_to_datafile("custom_ops/activations/_activation_ops.so"))
 
 
-@keras_utils.register_keras_custom_object
+@tf.keras.utils.register_keras_serializable(package='Addons')
 @tf.function
 def mish(x):
     """Mish: A Self Regularized Non-Monotonic Neural Activation Function.

tensorflow_addons/activations/rrelu.py

@@ -18,14 +18,13 @@
 from __future__ import print_function
 
 import tensorflow as tf
-from tensorflow_addons.utils import keras_utils
 from tensorflow_addons.utils.resource_loader import get_path_to_datafile
 
 _activation_ops_so = tf.load_op_library(
     get_path_to_datafile("custom_ops/activations/_activation_ops.so"))
 
 
-@keras_utils.register_keras_custom_object
+@tf.keras.utils.register_keras_serializable(package='Addons')
 @tf.function
 def rrelu(x, lower=0.125, upper=0.3333333333333333, training=None, seed=None):
     """rrelu function.

tensorflow_addons/activations/softshrink.py

@@ -18,14 +18,13 @@
 from __future__ import print_function
 
 import tensorflow as tf
-from tensorflow_addons.utils import keras_utils
 from tensorflow_addons.utils.resource_loader import get_path_to_datafile
 
 _activation_ops_so = tf.load_op_library(
     get_path_to_datafile("custom_ops/activations/_activation_ops.so"))
 
 
-@keras_utils.register_keras_custom_object
+@tf.keras.utils.register_keras_serializable(package='Addons')
 @tf.function
 def softshrink(x, lower=-0.5, upper=0.5):
     """Soft shrink function.

tensorflow_addons/activations/sparsemax.py

@@ -19,10 +19,8 @@
 
 import tensorflow as tf
 
-from tensorflow_addons.utils import keras_utils
 
-
-@keras_utils.register_keras_custom_object
+@tf.keras.utils.register_keras_serializable(package='Addons')
 @tf.function
 def sparsemax(logits, axis=-1):
     """Sparsemax activation function [1].

tensorflow_addons/activations/tanhshrink.py

@@ -18,14 +18,13 @@
 from __future__ import print_function
 
 import tensorflow as tf
-from tensorflow_addons.utils import keras_utils
 from tensorflow_addons.utils.resource_loader import get_path_to_datafile
 
 _activation_ops_so = tf.load_op_library(
     get_path_to_datafile("custom_ops/activations/_activation_ops.so"))
 
 
-@keras_utils.register_keras_custom_object
+@tf.keras.utils.register_keras_serializable(package='Addons')
 @tf.function
 def tanhshrink(x):
     """Applies the element-wise function: x - tanh(x)

tensorflow_addons/callbacks/README.md

@@ -16,8 +16,7 @@
 In order to conform with the current API standard, all callbacks
 must:
  * Inherit from `tf.keras.callbacks.Callback`.
- * [Register as a keras global object](https://github.com/tensorflow/addons/blob/master/tensorflow_addons/utils/keras_utils.py)
-  so it can be serialized properly.
+ * Register as a keras global object so it can be serialized properly: `@tf.keras.utils.register_keras_serializable(package='Addons')`
  * Add the addon to the `py_library` in this sub-package's BUILD file.
 
 #### Testing Requirements