typing activation functions (#975)

* typing activation functions

* update typing info

* sanity correction

tensorflow_addons/activations/hardshrink.py

@@ -14,13 +14,18 @@
 # ==============================================================================
 
 import tensorflow as tf
+from tensorflow_addons.utils.types import Number
+
+from tensorflow_addons.utils import types
 from tensorflow_addons.utils.resource_loader import LazySO
 
 _activation_so = LazySO("custom_ops/activations/_activation_ops.so")
 
 
 @tf.keras.utils.register_keras_serializable(package="Addons")
-def hardshrink(x, lower=-0.5, upper=0.5):
+def hardshrink(
+    x: types.TensorLike, lower: Number = -0.5, upper: Number = 0.5
+) -> tf.Tensor:
     """Hard shrink function.
 
     Computes hard shrink function:

tensorflow_addons/activations/lisht.py

@@ -14,13 +14,15 @@
 # ==============================================================================
 
 import tensorflow as tf
+
+from tensorflow_addons.utils import types
 from tensorflow_addons.utils.resource_loader import LazySO
 
 _activation_so = LazySO("custom_ops/activations/_activation_ops.so")
 
 
 @tf.keras.utils.register_keras_serializable(package="Addons")
-def lisht(x):
+def lisht(x: types.TensorLike) -> tf.Tensor:
     """LiSHT: Non-Parameteric Linearly Scaled Hyperbolic Tangent Activation Function.
 
     Computes linearly scaled hyperbolic tangent (LiSHT): `x * tanh(x)`

tensorflow_addons/activations/mish.py

@@ -14,13 +14,15 @@
 # ==============================================================================
 
 import tensorflow as tf
+
+from tensorflow_addons.utils import types
 from tensorflow_addons.utils.resource_loader import LazySO
 
 _activation_so = LazySO("custom_ops/activations/_activation_ops.so")
 
 
 @tf.keras.utils.register_keras_serializable(package="Addons")
-def mish(x):
+def mish(x: types.TensorLike) -> tf.Tensor:
     """Mish: A Self Regularized Non-Monotonic Neural Activation Function.
 
     Computes mish activation: x * tanh(softplus(x))

tensorflow_addons/activations/rrelu.py

@@ -14,10 +14,14 @@
 # ==============================================================================
 
 import tensorflow as tf
+from tensorflow_addons.utils.types import Number
+from tensorflow_addons.utils import types
+from typing import Optional
 
 
 @tf.keras.utils.register_keras_serializable(package='Addons')
-def rrelu(x, lower=0.125, upper=0.3333333333333333, training=None, seed=None):
+def rrelu(x: types.TensorLike, lower: Number = 0.125, upper: Number = 0.3333333333333333,
+          training: Optional[str] = None, seed: Optional[str] = None) -> tf.Tensor :
     """rrelu function.
 
     Computes rrelu function:

tensorflow_addons/activations/softshrink.py

@@ -14,13 +14,18 @@
 # ==============================================================================
 
 import tensorflow as tf
+from tensorflow_addons.utils.types import Number
+
+from tensorflow_addons.utils import types
 from tensorflow_addons.utils.resource_loader import LazySO
 
 _activation_so = LazySO("custom_ops/activations/_activation_ops.so")
 
 
 @tf.keras.utils.register_keras_serializable(package="Addons")
-def softshrink(x, lower=-0.5, upper=0.5):
+def softshrink(
+    x: types.TensorLike, lower: Number = -0.5, upper: Number = 0.5
+) -> tf.Tensor:
     """Soft shrink function.
 
     Computes soft shrink function:

tensorflow_addons/activations/sparsemax.py

@@ -15,10 +15,12 @@
 
 import tensorflow as tf
 
+from tensorflow_addons.utils import types
+
 
 @tf.keras.utils.register_keras_serializable(package="Addons")
 @tf.function
-def sparsemax(logits, axis=-1):
+def sparsemax(logits: types.TensorLike, axis: int = -1) -> tf.Tensor:
     """Sparsemax activation function [1].
 
     For each batch `i` and class `j` we have

tensorflow_addons/activations/tanhshrink.py

@@ -14,13 +14,15 @@
 # ==============================================================================
 
 import tensorflow as tf
+
+from tensorflow_addons.utils import types
 from tensorflow_addons.utils.resource_loader import LazySO
 
 _activation_so = LazySO("custom_ops/activations/_activation_ops.so")
 
 
 @tf.keras.utils.register_keras_serializable(package="Addons")
-def tanhshrink(x):
+def tanhshrink(x: types.TensorLike) -> tf.Tensor:
     """Applies the element-wise function: x - tanh(x)
 
     Args:

tools/ci_build/verify/check_typing_info.py

@@ -25,14 +25,6 @@
 # TODO: add types and remove all elements from
 # the exception list.
 EXCEPTION_LIST = [
-    tensorflow_addons.activations.hardshrink,
-    tensorflow_addons.activations.gelu,
-    tensorflow_addons.activations.lisht,
-    tensorflow_addons.activations.mish,
-    tensorflow_addons.activations.tanhshrink,
-    tensorflow_addons.activations.sparsemax,
-    tensorflow_addons.activations.softshrink,
-    tensorflow_addons.activations.rrelu,
     tensorflow_addons.callbacks.TQDMProgressBar,
     tensorflow_addons.callbacks.TimeStopping,
     tensorflow_addons.image.connected_components,