Typing the gelu activation function.

tensorflow_addons/activations/gelu.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 gelu(x, approximate=True):
+def gelu(x: types.TensorLike, approximate: bool = True) -> tf.Tensor:
     """Gaussian Error Linear Unit.
 
     Computes gaussian error linear:

tensorflow_addons/utils/types.py

@@ -1,26 +1,44 @@
-# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Types for typing functions signatures."""
-
-from typing import Union, Callable
-
-import numpy as np
-import tensorflow as tf
-
-Initializer = Union[None, dict, str, Callable]
-Regularizer = Union[None, dict, str, Callable]
-Constraint = Union[None, dict, str, Callable]
-
-FloatTensorLike = Union[tf.Tensor, float, np.float16, np.float32, np.float64]
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Types for typing functions signatures."""
+
+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]
+FloatTensorLike = Union[tf.Tensor, float, np.float16, np.float32, np.float64]