[Typing][B-52] Add type annotations for python/paddle/onnx/export.py (#66862)

enkilee · SigureMo · web-flow · commit 753329a9ff2e · 2024-08-01T18:33:37.000+08:00
---------

Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/jit/api.py b/python/paddle/jit/api.py
@@ -88,11 +88,12 @@
 )
 
 if TYPE_CHECKING:
+    from paddle import Tensor
     from paddle._typing import NestedStructure
     from paddle.static import InputSpec
 
     class _SaveOptions(TypedDict):
-        output_spec: NotRequired[Sequence[InputSpec]]
+        output_spec: NotRequired[Sequence[Tensor | int]]
         with_hook: NotRequired[bool]
         combine_params: NotRequired[bool]
         clip_extra: NotRequired[bool]
@@ -880,7 +881,7 @@ def __call__(
         self,
         layer: Layer | Callable[..., Any],
         path: str,
-        input_spec: Sequence[InputSpec | paddle.Tensor | object] | None = ...,
+        input_spec: Sequence[InputSpec | Tensor | object] | None = ...,
         **configs: Unpack[_SaveOptions],
     ) -> None:
         ...
@@ -891,7 +892,7 @@ def _run_save_pre_hooks(func: _SaveFunction) -> _SaveFunction:
     def wrapper(
         layer: Layer | Callable[..., Any],
         path: str,
-        input_spec: Sequence[InputSpec | paddle.Tensor | object] | None = None,
+        input_spec: Sequence[InputSpec | Tensor | object] | None = None,
         **configs: Unpack[_SaveOptions],
     ) -> None:
         global _save_pre_hooks
@@ -953,7 +954,7 @@ def _get_function_names_from_layer(layer: Layer) -> list[str]:
 def save(
     layer: Layer | Callable[..., Any],
     path: str,
-    input_spec: Sequence[InputSpec | paddle.Tensor | object] | None = None,
+    input_spec: Sequence[InputSpec | Tensor | object] | None = None,
     **configs: Unpack[_SaveOptions],
 ) -> None:
     """
diff --git a/python/paddle/onnx/export.py b/python/paddle/onnx/export.py
@@ -11,15 +11,32 @@
 # 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
 
 import os
+from typing import TYPE_CHECKING, Sequence
 
 from paddle.utils import try_import
 
+if TYPE_CHECKING:
+    from typing_extensions import Unpack
+
+    from paddle import Tensor
+    from paddle.jit.api import _SaveOptions
+    from paddle.nn import Layer
+    from paddle.static import InputSpec
+
+
 __all__ = []
 
 
-def export(layer, path, input_spec=None, opset_version=9, **configs):
+def export(
+    layer: Layer,
+    path: str,
+    input_spec: Sequence[InputSpec | Tensor | object] | None = None,
+    opset_version: int = 9,
+    **configs: Unpack[_SaveOptions],
+) -> None:
     """
     Export Layer to ONNX format, which can use for inference via onnxruntime or other backends.
     For more details, Please refer to `paddle2onnx <https://github.com/PaddlePaddle/paddle2onnx>`_ .
@@ -28,7 +45,7 @@ def export(layer, path, input_spec=None, opset_version=9, **configs):
         layer (Layer): The Layer to be exported.
         path (str): The path prefix to export model. The format is ``dirname/file_prefix`` or ``file_prefix`` ,
             and the exported ONNX file suffix is ``.onnx`` .
-        input_spec (list[InputSpec|Tensor], optional): Describes the input of the exported model's forward
+        input_spec (list[InputSpec|Tensor]|None, optional): Describes the input of the exported model's forward
             method, which can be described by InputSpec or example Tensor. If None, all input variables of
             the original Layer's forward method would be the inputs of the exported ``ONNX`` model. Default: None.
         opset_version(int, optional): Opset version of exported ONNX model.
@@ -62,7 +79,7 @@ def export(layer, path, input_spec=None, opset_version=9, **configs):
             ...     x_spec = paddle.static.InputSpec(shape=[None, 128], dtype='float32')
             ...     paddle.onnx.export(model, 'linear_net', input_spec=[x_spec])
             ...
-            >>> # doctest: +SKIP('raise ImportError')
+            >>> # doctest: +SKIP('Need install Paddle2ONNX')
             >>> export_linear_net()
 
             >>> class Logic(paddle.nn.Layer):
@@ -83,7 +100,7 @@ def export(layer, path, input_spec=None, opset_version=9, **configs):
             ...     # Static and run model.
             ...     paddle.jit.to_static(model)
             ...     out = model(x, y, z=True)
-            ...     paddle.onnx.export(model, 'pruned', input_spec=[x, y, True], output_spec=[out], input_names_after_prune=[x])
+            ...     paddle.onnx.export(model, 'pruned', input_spec=[x, y, True], output_spec=[out], input_names_after_prune=[x.name])
             ...
             >>> export_logic()
     """
