[Typing] Fix margin_cross_entropy argument reduction missing None type (#66007)

SigureMo · web-flow · commit d78c059da871 · 2024-07-14T22:27:56.000+08:00
* [Typing] Fix `margin_cross_entropy` arg reduction missing `None` type

* fix some typing

* ignore more
diff --git a/python/paddle/base/data_feeder.py b/python/paddle/base/data_feeder.py
@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import struct
+from typing import TYPE_CHECKING
 
 import numpy as np
 
 from paddle import pir
-from paddle._typing.dtype_like import DTypeLike
 
 from ..pir import Value
 from ..pir.core import _PADDLE_PIR_DTYPE_2_NUMPY_DTYPE, ParameterMeta
@@ -31,6 +33,10 @@
     in_pir_mode,
 )
 
+if TYPE_CHECKING:
+    from paddle._typing import DTypeLike
+    from paddle._typing.dtype_like import _DTypeLiteral
+
 __all__ = []
 
 _PADDLE_DTYPE_2_NUMPY_DTYPE = {
@@ -92,7 +98,7 @@ def convert_uint16_to_float(data):
     return np.reshape(new_data, data.shape)
 
 
-def convert_dtype(dtype: DTypeLike) -> str:
+def convert_dtype(dtype: DTypeLike) -> _DTypeLiteral:
     if isinstance(dtype, core.VarDesc.VarType):
         if dtype in _PADDLE_DTYPE_2_NUMPY_DTYPE:
             return _PADDLE_DTYPE_2_NUMPY_DTYPE[dtype]
diff --git a/python/paddle/distribution/lkj_cholesky.py b/python/paddle/distribution/lkj_cholesky.py
@@ -28,7 +28,7 @@
 
 if TYPE_CHECKING:
     from paddle import Tensor
-    from paddle._typing import DTypeLike
+    from paddle._typing.dtype_like import _DTypeLiteral
 
 
 __all__ = ["LKJCholesky"]
@@ -147,7 +147,7 @@ class LKJCholesky(distribution.Distribution):
     """
 
     concentration: Tensor
-    dtype: DTypeLike
+    dtype: _DTypeLiteral
     dim: int
     sample_method: Literal["onion", "cvine"]
 
diff --git a/python/paddle/io/dataloader/batch_sampler.py b/python/paddle/io/dataloader/batch_sampler.py
@@ -70,7 +70,7 @@ class BatchSampler(Sampler[Sequence[int]]):
 
             >>> np.random.seed(2023)
             >>> # init with dataset
-            >>> class RandomDataset(Dataset): # type: ignore[type-arg]
+            >>> class RandomDataset(Dataset):  # type: ignore[type-arg]
             ...     def __init__(self, num_samples):
             ...         self.num_samples = num_samples
             ...
@@ -226,7 +226,7 @@ class DistributedBatchSampler(BatchSampler):
             >>> from paddle.io import Dataset, DistributedBatchSampler
 
             >>> # init with dataset
-            >>> class RandomDataset(Dataset):
+            >>> class RandomDataset(Dataset):  # type: ignore[type-arg]
             ...     def __init__(self, num_samples):
             ...         self.num_samples = num_samples
             ...
@@ -386,7 +386,7 @@ def set_epoch(self, epoch: int) -> None:
                 >>> from paddle.io import Dataset, DistributedBatchSampler
 
                 >>> # init with dataset
-                >>> class RandomDataset(Dataset):
+                >>> class RandomDataset(Dataset):  # type: ignore[type-arg]
                 ...     def __init__(self, num_samples):
                 ...         self.num_samples = num_samples
                 ...
diff --git a/python/paddle/io/dataloader/sampler.py b/python/paddle/io/dataloader/sampler.py
@@ -69,7 +69,7 @@ class Sampler(Generic[_T]):
             >>> import numpy as np
             >>> from paddle.io import Dataset, Sampler
 
-            >>> class RandomDataset(Dataset):
+            >>> class RandomDataset(Dataset):  # type: ignore[type-arg]
             ...     def __init__(self, num_samples):
             ...         self.num_samples = num_samples
             ...
@@ -81,15 +81,15 @@ class Sampler(Generic[_T]):
             ...     def __len__(self):
             ...         return self.num_samples
             ...
-            >>> class MySampler(Sampler):
+            >>> class MySampler(Sampler):  # type: ignore[type-arg]
             ...     def __init__(self, data_source):
             ...         self.data_source = data_source
             ...
             ...     def __iter__(self):
-            ...         return iter(range(len(self.data_source)))
+            ...         return iter(range(len(self.data_source)))  # type: ignore[arg-type]
             ...
             ...     def __len__(self):
-            ...         return len(self.data_source)
+            ...         return len(self.data_source)  # type: ignore[arg-type]
             ...
             >>> sampler = MySampler(data_source=RandomDataset(100))
 
@@ -142,7 +142,7 @@ class SequenceSampler(Sampler[int]):
             >>> import numpy as np
             >>> from paddle.io import Dataset, SequenceSampler
 
-            >>> class RandomDataset(Dataset): # type: ignore[type-arg]
+            >>> class RandomDataset(Dataset):  # type: ignore[type-arg]
             ...     def __init__(self, num_samples):
             ...         self.num_samples = num_samples
             ...
@@ -205,7 +205,7 @@ class RandomSampler(Sampler[int]):
             >>> from paddle.io import Dataset, RandomSampler
 
             >>> np.random.seed(2023)
-            >>> class RandomDataset(Dataset):
+            >>> class RandomDataset(Dataset):  # type: ignore[type-arg]
             ...     def __init__(self, num_samples):
             ...         self.num_samples = num_samples
             ...
diff --git a/python/paddle/nn/functional/loss.py b/python/paddle/nn/functional/loss.py
@@ -2175,7 +2175,7 @@ def margin_cross_entropy(
     scale: float = ...,
     group=...,
     return_softmax: Literal[True] = ...,
-    reduction: _ReduceMode = ...,
+    reduction: _ReduceMode | None = ...,
 ) -> tuple[Tensor, Tensor]:
     ...
 
@@ -2190,7 +2190,7 @@ def margin_cross_entropy(
     scale: float = ...,
     group=...,
     return_softmax: Literal[False] = ...,
-    reduction: _ReduceMode = ...,
+    reduction: _ReduceMode | None = ...,
 ) -> Tensor:
     ...
 
@@ -2205,7 +2205,7 @@ def margin_cross_entropy(
     scale: float = ...,
     group=...,
     return_softmax: bool = ...,
-    reduction: _ReduceMode = ...,
+    reduction: _ReduceMode | None = ...,
 ) -> Tensor | tuple[Tensor, Tensor]:
     ...
 
