x-jPWBdZddlmZddlZddlZddlZddlmZddlm Z ddl m Z m Z ddl Ze rddlm ZgZGddeZGd d eZd#dZdZd$dZd%dZdZd$dZd%dZejaiadZd%dZej ddfd&d!Z!ej"ddfd'd"Z#dS)(z# Utilities of Auto SParsity (ASP). ) annotationsN)Enum) permutations) TYPE_CHECKINGAnyceZdZdZdZdZdZdS)MaskAlgoz A collection of all mask generating algorithms. There currently are three algorithms, `MASK_1D`, `MASK_2D_GREEDY` and `MASK_2D_BEST` get_mask_1dget_mask_2d_greedyget_mask_2d_bestN)__name__ __module__ __qualname____doc__MASK_1DMASK_2D_GREEDY MASK_2D_BESTY/var/www/html/banglarbhumi/venv/lib/python3.11/site-packages/paddle/incubate/asp/utils.pyr r $s) G)N%LLLrr c2eZdZdZdZdZed dZdS) CheckMethodzz A collection of all sparsity checking approaches. There currently are two methods, `CHECK_1D` and `CHECK_2D` check_mask_1d check_mask_2d mask_algor returnct|ts Jd|tjkr tjStjS)a Get sparsity checking method by mask generating algorithm. Args: mask_algo (MaskAlgo): The algorithm of mask generating. Returns: CheckMethod: The corresponded sparsity checking method. Examples: .. code-block:: python >>> import numpy as np >>> from paddle.incubate.asp import CheckMethod, MaskAlgo >>> print(CheckMethod.get_checking_method(MaskAlgo.MASK_1D)) CheckMethod.CHECK_1D >>> print(CheckMethod.get_checking_method(MaskAlgo.MASK_2D_GREEDY)) CheckMethod.CHECK_2D >>> print(CheckMethod.get_checking_method(MaskAlgo.MASK_2D_BEST)) CheckMethod.CHECK_2D z!mask_algo should be MaskAlgo type) isinstancer rrCHECK_1DCHECK_2D)rs rget_checking_methodzCheckMethod.get_checking_method8sM*)X..  /  . ( ( (' '' 'rN)rr rr)r rrrrr staticmethodr!rrrrr/sH HH(((\(((rrxnpt.NDArray[Any]rfloatc|}ttj|dj|jz S)a Return the density of the input tensor. Args: x (nparray): The input tensor. Returns: float, The density of :attr:`x`. Examples: .. code-block:: python >>> import paddle >>> import numpy as np >>> x = np.array([[0, 1, 3, 0], ... [1, 1, 0, 1]]) >>> out = paddle.incubate.asp.calculate_density(x) >>> print(out) 0.625 r)flattenr%npnonzerosize)r# x_flatteneds rcalculate_densityr,Vs90))++K K((+0 1 1K4D DDrct|jdks Jd|jd|z}|jd|zdkrgtj|jd|jd||z zf}||ddd|jdf<|j}|d||fS|d||jfS)a Reshape the input 2D matrix to shape (-1, m). If the second dimension of :attr:`mat` is not a multiples of :attr:`m`, then this function would pad the remainder with 0 before reshaping. .. math:: remainder = mat.shape[1] % m Args: mat (nparray): The input 2D matrix. m (int): The second dimension of reshaped matrix. Returns: tuple: A pair of the reshaped and padded matrix and the shape of padded matrix (non-reshaping). $The input mat should be a 2D matrix!rN)lenshaper(zerosreshape)matm remainder mat_paddedr3s r _reshape_1dr:rs sy>>Q    F    ! q I y|a!Xsy|SYq\Q]-KLMM (+ 111n ! n$% !!"a((%//{{2q!!39,,rr6nintr7boolct|jdkr3t|d|jd|\}}nt||\}}|D]+}t j|dj||z krdS,dS)a 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 to be a multiples of :attr:`m` if necessary. 1D `n:m` sparse pattern: At least :attr:`n` zeros in every :math:`1 \times m` block. Args: mat (nparray): The input matrix. n (int): n of `n:m` sparse pattern. m (int): m of `n:m` sparse pattern. Returns: bool: True if every row of :attr:`mat` is in 1D n:m sparse pattern, else False. Examples: .. code-block:: python >>> import numpy as np >>> import paddle.incubate.asp as sparsity >>> x = np.array([[0, 1, 3, 0], ... [1, 0, 0, 1]]) >>> y = sparsity.check_mask_1d(x, 2, 4) >>> print(y) True >>> x = np.array([[0, 1, 5, 4], ... [1, 0, 0, 1]]) >>> y = sparsity.check_mask_1d(x, 2, 4) >>> print(y) False >>> # x would be padded to shape (2, 8) >>> x = np.array([[0, 1, 0, 4, 6], ... [1, 0, 0, 1, 7]]) >>> y = sparsity.check_mask_1d(x, 2, 4) >>> print(y) True r0rFT)r2r3r:r5r(r)r*)r6r;r7 mat_flattenr3sub_mats rrrsN 39~~(Q ! )E)EqII UU(a00 U :g  q ! &!a% 0 055 1 4rct||\}}tj|}tj|}t|jdD]Q}||}tjtj|} d||| d|f<R||}|ddd|jdf|ddddf<|S)a 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` by zero to be a multiples of :attr:`m` before mask generation. 1D `n:m` sparse pattern: At least :attr:`n` zeros in every :math:`1 \times m` block. Args: mat (nparray): The input matrix. n (int): n of `n:m` sparse pattern. m (int): m of `n:m` sparse pattern. Returns: nparray: The 1D `n:m` sparse mask of :attr:`mat`. Examples: .. code-block:: python >>> import numpy as np >>> import paddle.incubate.asp as sparsity >>> mat = np.array([[0, 1, 5, 4], ... [2, 7, 3, 6]]) >>> mask = sparsity.get_mask_1d(mat, 2, 4) >>> print(mask) [[0 0 1 1] [0 1 0 1]] >>> y = sparsity.check_mask_1d(mask, 2, 4) >>> print(y) True rNr0) r:r( ones_likeranger3argsortabsolutetolistr5) r6r;r7r?r3 mask_flattenmaskir@min_order_indicess rr r s:%S!,,K< ,,L <  D ;$Q' ( (<<a.Jr{7';';<<:; Q)"1"-4466677''..Laaa39Q</0DAAAJ Krct|jdks Jd|jd|z}|jd|z}|dkr |jdn|jd||z z|dkr |jdn|jd||z zf}tj|}||d|jdd|jdf<tj|d||z}d}t d|jd|D]h}||z} t d|jd|D]D} | |z} tj||| | | fd} | ||<|dz }Ei||jfS)a3 Reshape the input 2D matrix to shape (-1, :math:`m \times m`). In each dimension of :attr:`mat`, if it is not a multiples of :attr:`m`, then this function would pad the remainder with 0 before reshaping. .. math:: remainder_0 = mat.shape[0] % m \\ remainder_1 = mat.shape[1] % m Args: mat (nparray): The input 2D matrix. m (int): The square root of second dimension of reshaped matrix. Returns: tuple: A pair of the reshaped and padded matrix and the shape of padded matrix (non-reshaping). r.r/rr0Nr1)r2r3r(r4emptyr5rCsqueeze) r6r7 remainder_0 remainder_1 new_shaper9r?curr_idx row_startrow_end col_startcol_endr@s r _reshape_2drVs" sy>>Q    F   )A,"K)A,"K$q(( ! cila+o.N#q(( ! cila+o.NI)$$J14J~1~~1~-.(9%%--b!a%88KH1j.q1155 a-q*"21"5q99  I!mGj9W,i.??@HHLLG%,K ! MHH    ( ((rc t||\}}|D]}tjtj|||dk}tjtj|d||z kdkr6tjtj|d||z kdkrdSdS)a 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 :attr:`m` if necessary. 2D `n:m` sparse pattern: At least :math:`n \times n` zeros in every :math:`m \times m` block under the constraint of at least :attr:`n` zeros for each row and column. Args: mat (nparray): The input matrix. n (int): n of `n:m` sparse pattern. m (int): m of `n:m` sparse pattern. Returns: bool: True if every :math:`m \times m` block of the input matrix :attr:`mat` is in 2D `n:m` sparse pattern, else False. Examples: .. code-block:: python >>> import numpy as np >>> import paddle.incubate.asp as sparsity >>> x = np.array([[0, 8, 9, 0], ... [9, 0, 0, 10], ... [5, 0, 0, 6], ... [0, 4, 6, 0]]) >>> y = sparsity.check_mask_2d(x, 2, 4) >>> print(y) True >>> x = np.array([[0, 8, 0, 9], ... [9, 0, 0, 10], ... [0, 5, 0, 6], ... [0, 4, 6, 0]]) >>> y = sparsity.check_mask_2d(x, 2, 4) >>> print(y) True >>> # x would be padded to shape (8, 8) >>> x = np.array([[0, 8, 0, 9], ... [9, 0, 7, 0], ... [0, 5, 0, 6], ... [3, 0, 6, 0], ... [1, 1, 0, 1]]) >>> y = sparsity.check_mask_2d(x, 2, 4) >>> print(y) True rr0axisFT)rVr(rErMr5sum)r6r;r7r9r3r@sub_masks rrrs^$C++J;rz'//!Q*?*?@@AAAE F26(+++q1u5 6 6! ; ; F26(+++q1u5 6 6! ; ;55 4rct|\}}tj|d}t t |D]!}tjtj||}tj||}tj|} fd| D} tj } tj } t t | dz ddD]n} | | }| |d|ks| |d|kr/d||d|df<| |dxxdz cc<| |dxxdz cc<o#tj |}d}t d|dD]<}|z}t d|dD]}|z}|||||||f<|dz }=|d|j dd|j dfS)a 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. 2D `n:m` sparse pattern: At least :math:`n \times n` zeros in every :math:`m \times m` block under the constraint of at least :attr:`n` zeros for each row and column. Greedily generating: For each :math:`m \times m` block, selecting values to keep in descent order. Args: mat (nparray): The input matrix. n (int): n of `n:m` sparse pattern. m (int): m of `n:m` sparse pattern. Returns: nparray: The 2D `n:m` sparse mask of :attr:`mat`. Examples: .. code-block:: python >>> import numpy as np >>> import paddle.incubate.asp as sparsity >>> mat = np.array([[9, 8, 3, 7], ... [9, 2, 1, 10], ... [5, 1, 3, 6], ... [2, 4, 6, 1]]) >>> mask = sparsity.get_mask_2d_greedy(mat, 2, 4) >>> print(mask) [[1. 1. 0. 0.] [1. 0. 0. 1.] [0. 0. 1. 1.] [0. 1. 1. 0.]] >>> y = sparsity.check_mask_2d(mask, 2, 4) >>> print(y) True r1c>g|]}t|z |zfSr)r<).0r#r7s r z&get_mask_2d_greedy..{s9   $%SQZZQ    rr0rg?N) rVr( zeros_liker5rCr2rErMrD collectionsCounterrLr3)r6r;r7r9r3 mask_paddedidxr@r[min_order_1d_indicesmin_order_2d_indices row_counter col_counterrI matrix_entryrHrQrRrSrTrUs ` rr r NsUJ$C++J- ++33B1==KS__%%..+bjC99:::k#.//!z'22    )=   ")++ !)++ s/0014b"== . .A/2LLO,11LO,1191a"899$U111X.qqq'')))(69%'')))rc t||}t||\}}tj|d||}tjtj|||jd||zjd}||dd|dd<tj |}d} td|d|D]<} | |z} td|d|D]} | |z} || || | | | f<| dz } =|d|jdd|jdfS)a 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 dimension of :attr:`mat` by zero to be a multiples of :attr:`m` before mask generation. 2D `n:m` sparse pattern: At least :math:`n \times n` zeros in every :math:`m \times m` block under the constraint of at least :attr:`n` zeros for each row and column. *Note*: L1 norm of sparse matrix from `Best` API is greater than or equal to the one from `Greedy`. Args: mat (nparray): The input matrix. n (int): n of `n:m` sparse pattern. m (int): m of `n:m` sparse pattern. Returns: nparray: The 1D `n:m` sparse mask of :attr:`mat`. Examples: .. code-block:: python >>> import numpy as np >>> import paddle.incubate.asp as sparsity >>> mat = np.array([[2, 8, 9, 9], ... [9, 1, 3, 9], ... [5, 6, 3, 9], ... [2, 4, 6, 9]]) >>> mask_greedy = sparsity.get_mask_2d_greedy(mat, 2, 4) >>> mask_best = sparsity.get_mask_2d_best(mat, 2, 4) >>> print("L1 norm of `greedy` sparse matrix", np.multiply(mat, mask_greedy).sum()) L1 norm of `greedy` sparse matrix 56.0 >>> print("L1 norm of `best` sparse matrix", np.multiply(mat, mask_best).sum()) L1 norm of `best` sparse matrix 61.0 r1rr0rXN) rwrVr(rBr5argmaxmatmulr3TrLrC)r6r;r7rtr?r3rGpmaxrHrQrRrSrTrUs rr r saD*!Q//H$S!,,K< ,,44RA>>L 9 +x//q0A1q5IIKLL    D tAAAw'LO 8E??DH1eAh** a-q%(A..  I!mG9Eh9OD7"Ig$55 6 MHH  #)A,#)A,. //rr.tensor func_namec|j}|j}|t}t |t sJdt |ttj t|j d}t|dkr| d|d}nXt|dkr$| |d|d}n!t|dkr,| |d|dz|d}nt|dkr|gd |d|dz|dz|d}|||| }| |d|d|d|dggd|Std t||||| }| ||S) aI 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. Args: tensor (nparray): The input tensor. func_name (MaskAlgo, optional): The function name to generate sparse mask. Default is `MaskAlgo.MASK_1D`. All options please refer to `MaskAlgo`. n (int, optional): n of `n:m` sparse pattern. Default is 2. m (int, optional): m of `n:m` sparse pattern. Default is 4. Returns: nparray: The `n:m` sparse mask of :attr:`tensor` generated by :attr:`func_name`. Examples: .. code-block:: python >>> import numpy as np >>> import paddle.incubate.asp as sparsity >>> tensor = np.array([[2, 8, 9, 9], ... [9, 1, 3, 9], ... [5, 6, 3, 9], ... [2, 4, 6, 9]]) >>> mask_1d = sparsity.create_mask(tensor, func_name=sparsity.MaskAlgo.MASK_1D) >>> print(mask_1d) [[0 0 1 1] [1 0 0 1] [0 1 0 1] [0 0 1 1]] >>> mask_2d = sparsity.create_mask(tensor, func_name=sparsity.MaskAlgo.MASK_2D_BEST) >>> print(mask_2d) [[0 1 1 0] [1 0 0 1] [1 1 0 0] [0 0 1 1]] zMfunc_name argument of create_mask is only accepted as type MaskAlgo. But got Nr0rr.r}rr0rr.r;r7gThe dimension of input tensor is not supported in create_mask, Only dimension < 4 is supported but got )r3dtypeastyper%rr typegetattrsysmodulesr valuer2r5 transpose ValueError) r~rr;r7r3rtfuncrHs r create_maskrs P LE LE eA i * * % ?? % % * 3;x()/4 @ @D 5zzQ IIaq " " Uq IIeAha ) ) Uq IIeAhq)58 4 4 Uq KK % % - - !HuQx %( *E!H  tAa    LL%(E!HeAhaA B B Y||| $ $ VE]]  D7:5zz D D   4Q!   D <<   % %e , ,,rc^|j}|t}t|tksJdt|t t jt|j d}t|dkr| d|d}nt|dkr#| |d|d}nt|dkr,| |d|dz|d}n}t|dkrK| gd |d|dz|dz|dg}ntd t||||| S) a 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. Args: tensor (nparray): The input tensor. func_name (CheckMethod, optional): The function name to generate sparse mask. Default is `CheckMethod.CHECK_1D`. All options please refer to `CheckMethod`. n (int, optional): n of `n:m` sparse pattern. Default is 2. m (int, optional): m of `n:m` sparse pattern. Default is 4. Returns: bool: True if tensor pass checking of function given by :attr:`func_name`, else False. Examples: .. code-block:: python >>> import numpy as np >>> import paddle.incubate.asp as sparsity >>> tensor = np.array([[2, 8, 9, 9], ... [9, 1, 3, 9], ... [5, 6, 3, 9], ... [2, 4, 6, 9]]) >>> mask_1d = sparsity.create_mask(tensor, func_name=sparsity.MaskAlgo.MASK_1D) >>> print(mask_1d) [[0 0 1 1] [1 0 0 1] [0 1 0 1] [0 0 1 1]] >>> y = sparsity.check_sparsity(mask_1d, func_name=sparsity.CheckMethod.CHECK_1D) >>> print(y) True >>> y = sparsity.check_sparsity(mask_1d, func_name=sparsity.CheckMethod.CHECK_2D) >>> print(y) True zSfunc_name argument of check_sparsity is only accepted as type CheckMethod. But got Nr0rr.rr}rrr)r3rr%rrrrrr rr2r5rr)r~rr;r7r3rrs rcheck_sparsityrHsP LE eA  ??k ) ) ) % ?? % % * ) ) 3;x()/4 @ @D 5zzQ IIaq " " Uq IIeAha ) ) Uq IIeAhq)58 4 4 Uq KK % % - - 1Xa 58 +U1X 6   D7:5zz D D   4Q!   r)r#r$rr%)r6r$r;r<r7r<rr=)r6r$r;r<r7r<rr$) r~r$rr r;r<r7r<rr$) r~r$rrr;r<r7r<rr=)$r __future__rrar threadingenumr itertoolsrtypingrrnumpyr( numpy.typingnpt__all__r rr,r:rr rVrr Lockrprlrwr rrrrrrrrs1#""""" """"""%%%%%%%% &&&&&t&&&$($($($($($$($($(NEEEE8---8////d''''T()()()V6666rF0F0F0F0R)).**&&&R50505050t#*   I-I-I-I-I-\)1   AAAAAAAr