[Typing] 修改 uniform.py 中输入参数的类型 (#65742)

megemini · web-flow · commit f3d8738f1be7 · 2024-07-06T20:29:13.000+08:00
diff --git a/python/paddle/distribution/uniform.py b/python/paddle/distribution/uniform.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING, Sequence, Union
 
 import numpy as np
 import numpy.typing as npt
@@ -27,7 +27,18 @@
 from paddle.tensor import random
 
 if TYPE_CHECKING:
+    from typing_extensions import TypeAlias
+
     from paddle import Tensor
+    from paddle._typing import NestedSequence
+
+    _UniformBoundary: TypeAlias = Union[
+        float,
+        Sequence[float],
+        NestedSequence[float],
+        npt.NDArray[Union[np.float32, np.float64]],
+        Tensor,
+    ]
 
 
 class Uniform(distribution.Distribution):
@@ -106,23 +117,14 @@ class Uniform(distribution.Distribution):
             Tensor(shape=[1], dtype=float32, place=Place(cpu), stop_gradient=True,
                 [0.50000000])
     """
+
     low: Tensor
     high: Tensor
 
     def __init__(
         self,
-        low: (
-            float
-            | Sequence[float]
-            | npt.NDArray[np.float32 | np.float64]
-            | Tensor
-        ),
-        high: (
-            float
-            | Sequence[float]
-            | npt.NDArray[np.float32 | np.float64]
-            | Tensor
-        ),
+        low: _UniformBoundary,
+        high: _UniformBoundary,
         name: str | None = None,
     ) -> None:
         if not in_dynamic_mode():
