From 1a9d582b1521c6f108298684414e67e2ed615b6b Mon Sep 17 00:00:00 2001 From: haoyu <3253077866@qq.com> Date: Mon, 19 Aug 2024 20:28:53 +0800 Subject: [PATCH 1/2] update fused_rotary_position_embedding.py --- .../fused_rotary_position_embedding.py | 26 ++++++++++++------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/python/paddle/incubate/nn/functional/fused_rotary_position_embedding.py b/python/paddle/incubate/nn/functional/fused_rotary_position_embedding.py index b94d37acd7b4dc..de12af4281a869 100644 --- a/python/paddle/incubate/nn/functional/fused_rotary_position_embedding.py +++ b/python/paddle/incubate/nn/functional/fused_rotary_position_embedding.py @@ -12,23 +12,29 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + +from typing import TYPE_CHECKING from paddle import _C_ops from paddle.base.layer_helper import LayerHelper from paddle.framework import in_dynamic_or_pir_mode +if TYPE_CHECKING: + from paddle import Tensor + def fused_rotary_position_embedding( - q, - k=None, - v=None, - sin=None, - cos=None, - position_ids=None, - use_neox_rotary_style=True, - time_major=False, - rotary_emb_base=10000.0, -): + q: Tensor, + k: Tensor | None = None, + v: Tensor | None = None, + sin: Tensor | None = None, + cos: Tensor | None = None, + position_ids: Tensor | None = None, + use_neox_rotary_style: bool = True, + time_major: bool = False, + rotary_emb_base: float = 10000.0, +) -> Tensor: r""" Fused rotary position embedding. From b1074e71f6d6de175fd2647cde0f6d2992f63f22 Mon Sep 17 00:00:00 2001 From: "Lingchen.Yang" <97349759+haoyu2022@users.noreply.github.com> Date: Tue, 20 Aug 2024 16:47:12 +0800 Subject: [PATCH 2/2] Update python/paddle/incubate/nn/functional/fused_rotary_position_embedding.py Co-authored-by: megemini --- .../incubate/nn/functional/fused_rotary_position_embedding.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/paddle/incubate/nn/functional/fused_rotary_position_embedding.py b/python/paddle/incubate/nn/functional/fused_rotary_position_embedding.py index de12af4281a869..b99296d2dabdde 100644 --- a/python/paddle/incubate/nn/functional/fused_rotary_position_embedding.py +++ b/python/paddle/incubate/nn/functional/fused_rotary_position_embedding.py @@ -34,7 +34,7 @@ def fused_rotary_position_embedding( use_neox_rotary_style: bool = True, time_major: bool = False, rotary_emb_base: float = 10000.0, -) -> Tensor: +) -> tuple[Tensor, Tensor, Tensor]: r""" Fused rotary position embedding.