[Typing][C-74] Add type annotations for python/paddle/incubate/nn/functional/fused_matmul_bias.py (#66656)

Betelgeu · web-flow · commit 03e5e1af9adc · 2024-07-28T19:46:52.000+08:00
diff --git a/python/paddle/incubate/nn/functional/fused_matmul_bias.py b/python/paddle/incubate/nn/functional/fused_matmul_bias.py
@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
 from paddle import _C_ops, _legacy_C_ops
 from paddle.base.layer_helper import LayerHelper
 from paddle.framework import (
@@ -20,10 +24,18 @@
 )
 from paddle.tensor.linalg import matmul
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
 
 def fused_matmul_bias(
-    x, y, bias=None, transpose_x=False, transpose_y=False, name=None
-):
+    x: Tensor,
+    y: Tensor,
+    bias: Tensor | None = None,
+    transpose_x: bool = False,
+    transpose_y: bool = False,
+    name: str | None = None,
+) -> Tensor:
     """
     Applies matrix multiplication of two tensors and then bias addition if provided.
     This method requires CUDA version >= 11.6.
@@ -80,7 +92,13 @@ def fused_matmul_bias(
     return out
 
 
-def fused_linear(x, weight, bias=None, transpose_weight=False, name=None):
+def fused_linear(
+    x: Tensor,
+    weight: Tensor,
+    bias: Tensor | None = None,
+    transpose_weight: bool = False,
+    name: str | None = None,
+) -> Tensor:
     """
     Fully-connected linear transformation operator. This method requires CUDA version >= 11.6.
 
@@ -116,8 +134,13 @@ def fused_linear(x, weight, bias=None, transpose_weight=False, name=None):
 
 
 def fused_linear_activation(
-    x, y, bias, trans_x=False, trans_y=False, activation=None
-):
+    x: Tensor,
+    y: Tensor,
+    bias: Tensor,
+    trans_x: bool = False,
+    trans_y: bool = False,
+    activation: Literal['gelu', 'relu'] | None = None,
+) -> Tensor:
     """
     Fully-connected linear and activation transformation operator. This method requires CUDA version >= 11.6.
 
