[Typing][C-58] Add type annotations for python/paddle/incubate/asp/utils.py (#66012)

megemini · web-flow · commit dd6b46793c28 · 2024-07-14T17:08:08.000+08:00
diff --git a/python/paddle/incubate/asp/utils.py b/python/paddle/incubate/asp/utils.py
@@ -16,14 +16,20 @@
 Utilities of Auto SParsity (ASP).
 """
 
+from __future__ import annotations
+
 import collections
 import sys
 import threading
 from enum import Enum
 from itertools import permutations
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
+if TYPE_CHECKING:
+    import numpy.typing as npt
+
 __all__ = []
 
 
@@ -32,6 +38,7 @@ class MaskAlgo(Enum):
     A collection of all mask generating algorithms.
     There currently are three algorithms, `MASK_1D`, `MASK_2D_GREEDY` and `MASK_2D_BEST`
     """
+
     MASK_1D = 'get_mask_1d'
     MASK_2D_GREEDY = 'get_mask_2d_greedy'
     MASK_2D_BEST = 'get_mask_2d_best'
@@ -42,11 +49,12 @@ class CheckMethod(Enum):
     A collection of all sparsity checking approaches.
     There currently are two methods, `CHECK_1D` and `CHECK_2D`
     """
+
     CHECK_1D = 'check_mask_1d'
     CHECK_2D = 'check_mask_2d'
 
     @staticmethod
-    def get_checking_method(mask_algo):
+    def get_checking_method(mask_algo: MaskAlgo) -> CheckMethod:
         r"""
         Get sparsity checking method by mask generating algorithm.
 
@@ -75,7 +83,7 @@ def get_checking_method(mask_algo):
             return CheckMethod.CHECK_2D
 
 
-def calculate_density(x):
+def calculate_density(x: npt.NDArray[Any]) -> float:
     r"""
 
     Return the density of the input tensor.
@@ -131,7 +139,7 @@ def _reshape_1d(mat, m):
         return mat.reshape(-1, m), mat.shape
 
 
-def check_mask_1d(mat, n, m):
+def check_mask_1d(mat: npt.NDArray[Any], n: int, m: int) -> bool:
     r"""
     Check if every row of the input matrix :attr:`mat` is in 1D `n:m` sparse pattern.
     This function would pad the second dimension of :attr:`mat` by zero
@@ -181,7 +189,7 @@ def check_mask_1d(mat, n, m):
     return True
 
 
-def get_mask_1d(mat, n, m):
+def get_mask_1d(mat: npt.NDArray[Any], n: int, m: int) -> npt.NDArray[Any]:
     r"""
     Generate 1D `n:m` sparse pattern mask of the input matrix :attr:`mat`
     in row-directory. This function would pad the second dimension of :attr:`mat`
@@ -266,7 +274,7 @@ def _reshape_2d(mat, m):
     return mat_flatten, mat_padded.shape
 
 
-def check_mask_2d(mat, n, m):
+def check_mask_2d(mat: npt.NDArray[Any], n: int, m: int) -> bool:
     r"""
     Check if every :math:`m \times m` block of the input matrix :attr:`mat` is in 2D `n:m` sparse pattern.
     This function would pad each dimension of :attr:`mat` by zero to be a multiples of
@@ -323,7 +331,9 @@ def check_mask_2d(mat, n, m):
     return True
 
 
-def get_mask_2d_greedy(mat, n, m):
+def get_mask_2d_greedy(
+    mat: npt.NDArray[Any], n: int, m: int
+) -> npt.NDArray[Any]:
     r"""
     Greedily generate 2D `n:m` sparse pattern mask of the input matrix :attr:`mat`.
     This function would pad each dimension of :attr:`mat` by zero to be a multiples of :attr:`m` before mask generation.
@@ -439,7 +449,7 @@ def _compute_valid_2d_patterns(n, m):
         return valid_patterns
 
 
-def get_mask_2d_best(mat, n, m):
+def get_mask_2d_best(mat: npt.NDArray[Any], n: int, m: int) -> npt.NDArray[Any]:
     r"""
     Generate 2D `n:m` sparse pattern mask of the input matrix :attr:`mat`
     to form sparse matrix with maximum L1 norm .This function would pad each
@@ -495,7 +505,12 @@ def get_mask_2d_best(mat, n, m):
     return mask[: mat.shape[0], : mat.shape[1]]
 
 
-def create_mask(tensor, func_name=MaskAlgo.MASK_1D, n=2, m=4):
+def create_mask(
+    tensor: npt.NDArray[Any],
+    func_name: MaskAlgo = MaskAlgo.MASK_1D,
+    n: int = 2,
+    m: int = 4,
+) -> npt.NDArray[Any]:
     r"""
     Create `n:m` sparse pattern mask of the input tensor via function given by :attr:`func_name`.
     Currently only support tensor with dimension less than or equal to 4.
@@ -566,7 +581,12 @@ def create_mask(tensor, func_name=MaskAlgo.MASK_1D, n=2, m=4):
     return mask.reshape(shape).astype(dtype)
 
 
-def check_sparsity(tensor, func_name=CheckMethod.CHECK_1D, n=2, m=4):
+def check_sparsity(
+    tensor: npt.NDArray[Any],
+    func_name: CheckMethod = CheckMethod.CHECK_1D,
+    n: int = 2,
+    m: int = 4,
+) -> bool:
     r"""
     Check if input tensor is in `n:m` sparse pattern via function given by :attr:`func_name`.
     Currently only support tensor with dimension less than or equal to 4.
