diff --git a/python/paddle/incubate/nn/functional/fused_matmul_bias.py b/python/paddle/incubate/nn/functional/fused_matmul_bias.py index 56d5e30a506ab1..e34fa4e71608da 100644 --- 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.