[Typing][A-60] Add type annotations for paddle/nn/functional/input.py (#65317)

sunzhongkai588 · web-flow · commit e4cda9b244cd · 2024-06-22T00:50:43.000+08:00
diff --git a/python/paddle/nn/functional/input.py b/python/paddle/nn/functional/input.py
@@ -11,6 +11,9 @@
 # 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.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 from paddle import _C_ops
 
@@ -19,10 +22,17 @@
 from ...common_ops_import import Variable
 from ...framework import in_dynamic_or_pir_mode
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
 __all__ = []
 
 
-def one_hot(x, num_classes, name=None):
+def one_hot(
+    x: Tensor,
+    num_classes: int,
+    name: str | None = None,
+) -> Tensor:
     """
 
     The operator converts each id in the input `x` to an one-hot vector with a
@@ -70,7 +80,7 @@ def one_hot(x, num_classes, name=None):
            None by default.
 
     Returns:
-        Tensor: The one-hot representations of `x`. A Tensor with type float32.
+        Tensor, The one-hot representations of `x`. A Tensor with type float32.
 
     Examples:
         .. code-block:: python
@@ -117,7 +127,13 @@ def one_hot(x, num_classes, name=None):
         return one_hot_out
 
 
-def embedding(x, weight, padding_idx=None, sparse=False, name=None):
+def embedding(
+    x: Tensor,
+    weight: Tensor,
+    padding_idx: int | None = None,
+    sparse: bool = False,
+    name: str | None = None,
+) -> Tensor:
     r"""
     Used to lookup embeddings vector of ids provided by :attr:`x` .
 
@@ -157,7 +173,7 @@ def embedding(x, weight, padding_idx=None, sparse=False, name=None):
             True because sparse update is faster. But some optimizers does not support sparse update,
             such as :ref:`api_paddle_optimizer_adadelta_Adadelta` , :ref:`api_paddle_optimizer_adamax_Adamax` , :ref:`api_paddle_optimizer_lamb_Lamb`.
             In these cases, sparse must be False. Default: False.
-        padding_idx(int|long|None, optional): padding_idx needs to be in the interval [-weight.shape[0], weight.shape[0]).
+        padding_idx(int|None, optional): padding_idx needs to be in the interval [-weight.shape[0], weight.shape[0]).
             If :math:`padding\_idx < 0`, the :math:`padding\_idx` will automatically be converted
             to :math:`weight.shape[0] + padding\_idx` . It will output all-zero padding data whenever lookup
             encounters :math:`padding\_idx` in id. And the padding data will not be updated while training.
@@ -167,7 +183,7 @@ def embedding(x, weight, padding_idx=None, sparse=False, name=None):
            None by default.
 
     Returns:
-        Tensor: Embedding Tensor  mapped by x. The data type is the same as :attr:`weight`.
+        Tensor, Embedding Tensor  mapped by x. The data type is the same as :attr:`weight`.
 
     Examples:
 
