Typing the gelu activation function.

tensorflow_addons/activations/gelu.py

@@ -14,13 +14,18 @@
 # ==============================================================================
 
 import tensorflow as tf
+from typeguard import typechecked
+
+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 gelu(x, approximate=True):
+@tf.autograph.experimental.do_not_convert
+@typechecked
+def gelu(x: types.TensorLike, approximate: bool = True) -> tf.Tensor:
     """Gaussian Error Linear Unit.
 
     Computes gaussian error linear:

tensorflow_addons/utils/types.py

@@ -14,8 +14,28 @@
 # ==============================================================================
 """Types for typing functions signatures."""
 
-from typing import Union, Callable
+from typing import Union, Callable, List
+
+import numpy as np
+import tensorflow as tf
+
+
+Number = Union[float,
+               int,
+               np.float16,
+               np.float32,
+               np.float64,
+               np.int8,
+               np.int16,
+               np.int32,
+               np.int64,
+               np.uint8,
+               np.uint16,
+               np.uint32,
+               np.uint64]
 
 Initializer = Union[None, dict, str, Callable]
 Regularizer = Union[None, dict, str, Callable]
 Constraint = Union[None, dict, str, Callable]
+
+TensorLike = Union[List[Union[Number, list]], Number, np.ndarray, tf.Tensor]