[Typing][C-63] Add type annotations for python/paddle/incubate/framework/random.py (#66028)

megemini · SigureMo · web-flow · commit 219ad27a8554 · 2024-07-15T10:06:33.000+08:00
---------

Co-authored-by: Nyakku Shigure &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/incubate/framework/random.py b/python/paddle/incubate/framework/random.py
@@ -12,14 +12,46 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# TODO: define random api
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal, Protocol, Sequence, overload
+
 import paddle
 from paddle import base
 from paddle.base import core
 
+if TYPE_CHECKING:
+    from paddle._typing import PlaceLike
+
+    class _GeneratorState(Protocol):
+        def current_seed(self) -> int:
+            ...
+
+
 __all__ = []
 
 
+@overload
+def get_rng_state(
+    device: PlaceLike | None = ..., use_index: Literal[True] = ...
+) -> list[int]:
+    ...
+
+
+@overload
+def get_rng_state(
+    device: PlaceLike | None = ..., use_index: Literal[False] = ...
+) -> list[_GeneratorState]:
+    ...
+
+
+@overload
+def get_rng_state(
+    device: PlaceLike | None = ..., use_index: bool = ...
+) -> list[int] | list[_GeneratorState]:
+    ...
+
+
 def get_rng_state(device=None, use_index=False):
     """
     Get all random states of random generators of specified device.
@@ -79,7 +111,38 @@ def get_state(generator):
     return state_list
 
 
-def set_rng_state(state_list, device=None, use_index=False):
+@overload
+def set_rng_state(
+    state_list: Sequence[int],
+    device: PlaceLike | None = ...,
+    use_index: Literal[True] = ...,
+) -> None:
+    ...
+
+
+@overload
+def set_rng_state(
+    state_list: Sequence[_GeneratorState],
+    device: PlaceLike | None = ...,
+    use_index: Literal[False] = ...,
+) -> None:
+    ...
+
+
+@overload
+def set_rng_state(
+    state_list: Sequence[int] | Sequence[_GeneratorState],
+    device: PlaceLike | None = ...,
+    use_index: bool = ...,
+) -> None:
+    ...
+
+
+def set_rng_state(
+    state_list,
+    device=None,
+    use_index=False,
+):
     """
 
     Sets generator state for all device generators.
@@ -156,7 +219,10 @@ def set_state(generator, state):
         )
 
 
-def register_rng_state_as_index(state_list=None, device=None):
+def register_rng_state_as_index(
+    state_list: Sequence[_GeneratorState] | None = None,
+    device: PlaceLike | None = None,
+) -> list[int]:
     """
 
     The register_rng_state_as_index function creates and registers a new generator state within the generator.
