
    `i9              	       h   d dl mZ d dlmZ d dlZd dlZd dlZd dlm	Z	 d dlm
Z
 d dlmZ d dlmZ d dlmZ d d	lmZ d d
lmZmZ ej        ej        ej        fZ ed ej        dej        dej        d i          Zd Zd Zd Z	 d Z	 ej        	 	 d d            Z	 ej        d!d            Z 	 ej        dddej!        fd            Z"ej        d"d            Z#	 ej        	 	 	 	 	 d#d            Z$ej%        j&        Z'ej(        ej)                 j*        Z+de+d<   e'dk    re'dk     r
ej,        e+d<   d Z-dS )$    )defaultdict)IterableN)reduce)assume)settings)
strategies)numpy)SearchStrategy)_calculate_dynamic_qparams&_calculate_dynamic_per_channel_qparamsc                      d S )N r       |/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/torch/testing/_internal/hypothesis_utils.py<lambda>r       s    4 r   c                 ~   | \  }}}dt          j        t           j                  j        z   }t          j        t           j                  }|j        |z  |j        |z  }}t          ||z
  |z  ||z  |z             }t          ||z
  |z  ||z  |z             }	t          j	        |          t          j	        |	          fS N   )
torchfinfofloatepsiinfolongminmaxnpfloat32)
qparamsscale
zero_point_quantized_type
adjustment_long_type_infolong_minlong_max	min_value	max_values
             r   _get_valid_min_maxr)   &   s    )0&E:U[--11Jk%*--O(,z9?;NQ[;[hHX
*e3h6F6SUUIX
*e3h6F6SUUI:i  "*Y"7"777r   c                     d|v rxt           j        j        dk     rbd|v r	|d          sd|vod|v r	|d          pd|v}t          |           dk    od|vod|v}|r|r
|d         dk    rMt	          j        t          j                  j        |d<   t	          j        t          j                  j        |d<   n|d         d	k    rMt	          j        t          j	                  j        |d<   t	          j        t          j	                  j        |d<   nX|d         d
k    rLt	          j        t          j
                  j        |d<   t	          j        t          j
                  j        |d<   |                    d           t          j        | i |S )Nwidth)   C   r   	allow_nanallow_infinityr   r'   r(          @   )
hypothesisversion__version_info__lenr   r   float16r   r   r   float64popstfloats)argskwargsno_nan_and_infmin_and_max_not_specifieds       r   _floats_wrapperr@   3   s   &Z/@:MM V#?F;,?(? '&-&(I8H1I-I ,V+	 	 IIN &v%&v% 	"
  		E7 		Eg"$$&+k%-&@&@&D{#&+k%-&@&@&D{##B&&&+k%-&@&@&D{#&+k%-&@&@&D{##B&&&+k%-&@&@&D{#&+k%-&@&@&D{#

79d%f%%%r   c                  ,    d|vrd|d<   t          | i |S )Nr+   r1   )r@   )r<   r=   s     r   r;   r;   N   s)    fwD+F+++r   c                     t          |          \  }}t          |                                 |k               t          |                                 |k               dS )NT)r)   r   r   r   )tensorr   r'   r(   s       r   assume_not_overflowingrD   b   sO    -g66Iy
6::<<9$%%%
6::<<9$%%%4r   c                    |t           }t          |t          t          f          s|f} | t	          j        |                    }t          j        |          }|j        |j	        }	}t          |         }
|
|
}n+||n|}||	n|} | t	          j        ||                    }|#t          j        t          j                  j        }|#t          j        t          j                  j	        } | t          ||d                    }|||fS )N)r'   r(   r1   )r'   r(   r+   )_ALL_QINT_TYPES
isinstancelisttupler:   sampled_fromr   r   r   r   _ENFORCED_ZERO_POINTintegersr   r   r   r;   )drawdtypes	scale_min	scale_maxzero_point_minzero_point_maxquantized_type
_type_infoqminqmax_zp_enforcedr!   _zp_min_zp_maxr    s                  r   r   r   w   s    ~ ftUm,, T"/&1122N^,,J$D (7L!

(0$$n(0$$nT"+7KKKLL
K,,0	K,,0	D)yKKKLLE*n,,r   r   c                 *   |dk     sJ |t          |dz   d          }|dk     sJ ||dz   }t          j        t          j        ||          ||          }|                    fd          } | |                    t                              S )z8Return a strategy for array shapes (tuples of int >= 1).r1   N      )min_sizemax_sizec                 B    t          t          j        | d          k    S r   )r   int__mul__)x	max_numels    r   r   zarray_shapes.<locals>.<lambda>   s    vck1a/H/HI/U r   )r   r:   listsrL   filtermaprI   )rM   min_dimsmax_dimsmin_sidemax_siderc   	candidates        ` r   array_shapesrl      s     b====x!|R((b====a<Xx888V^___I$$%U%U%U%UVV	4	e$$%%%r   c                    t          |t                    r | |          }n | t          j        |                    }||t	          dddd          } | t          j        |||                    }t          t          j	        |          
                                p%t          j        |          
                                            |d fS  | |          }|&t          |          \  }}t	          ||ddd          } | t          j        |||                    }t          ||d                   \  }	}
t                              |d         d           }||}
||	|
|d         ffS )	N    .    .AFr1   r.   r+   dtypeelementsshaper/   r.   r+   r[   )rG   r
   r:   rJ   r;   stnparraysr   r   isnananyisinfr)   r   rK   get)rM   shapesrs   r   rr   _shapeXr'   r(   r    zpenforced_zps               r   rC   rC      su   &.)) /fbof--..dC5CCCHD586JJJKKBHQKKOO%%:!):):;<<<$wd7mmG1'::	9)Yu$)5 5 5T[uxvFFFGGA*1gaj99IE2&**71:t<<Kub'!*%%%r   c                    t          |t                    r | |          }n | t          j        |                    }||t	          dddd          } | t          j        t          j        ||                    }t          t          j
        |                                          p%t          j        |                                                      |d fS  | |          }|&t          |          \  }}t	          ||ddd          } | t          j        t          j        ||                    }t          ||d                   \  }}	t                              |d         d           }
|
|
}	t#          t          j                            d	|j        d
                    }t          j        |j                  }||d	<   d	||<   t          j        ||          }|||	||d         ffS )Nrn   ro   Fr1   rp   rq   ru   r[   r   r   )rG   r
   r:   rJ   r;   rv   rw   r   r   r   rx   ry   rz   r)   r   rK   r{   r`   randomrandintndimarange	transpose)rM   r|   rs   r   r}   r~   r'   r(   r    r   r   axispermute_axess                r   per_channel_tensorr      s   &.)) /fbof--..dC5CCCHD2:OOOPPBHQKKOO%%:!):):;<<<$wd7mmG1'::	9)Yu$)5 5 5T[rzHFKKKLLA6q'!*EEIE2&**71:t<<Kry  AFA..//D9QV$$LLOL
Q%%Aub$
+++r   r[   r      r,            Fc                       t          j        |           }  t          j        |           }  t          j        |           }  t          j        d|                    }||z  }||z  }t          |t                    r  t          j        |                    } fdt          |          D             } fdt          |          D             }d}||ft          |          z   }|}|r4  t          j                              }|r||ft          |          z   }|}|
@t          |
t          t          f          rt          |
          dk    s
J d            n|
gdz  }
  t          ||ft          |          z   f|	|
d                             }  t          |f|	|
d                             }  t          |f|	|
d	                             }|||||fS )
Nr   c                 >    g | ]} t          j                   S r   r:   rL   ).0_rM   feature_map_ranges     r   
<listcomp>ztensor_conv.<locals>.<listcomp>?  s+    [[[1bk+<=>>[[[r   c                 >    g | ]} t          j                   S r   r   )r   r   rM   kernel_ranges     r   r   ztensor_conv.<locals>.<listcomp>A  s*    LLLAttBK.//LLLr   Fr,   zNeed 3 qparams for X, w, br   )r|   rs   r   r[   )r:   rL   rG   r   rJ   rangerI   booleansrH   r6   rC   )rM   spatial_dimbatch_size_rangeinput_channels_per_group_rangeoutput_channels_per_group_ranger   r   
max_groupscan_be_transposedrs   r   
batch_sizeinput_channels_per_groupoutput_channels_per_groupgroupsinput_channelsoutput_channelsfeature_map_shapekernelstrweight_shape
bias_shaper~   Wbs   `    ``                  r   tensor_convr   )  sa    bk#3455J#t
34 6  6 $
45!7 !7T"+a,,--F-6N/&8O+x(( 9d2?;7788[[[[[kHZHZ[[[LLLLL{9K9KLLLG	B#%=>wOL J )T"+--   	)*,EFwWL(J ge}-- 	$w<<1$$$&B$$$$i!mGV	^$u->'?'??B71:/ / / 	0 	0A 	VL?X#AJ( ( ( 	) 	)AVJ=8#AJ( ( ( 	) 	)A aFBr   deadline)r,   r0   r   )r\   r   r   timeoutc                      t           dk     r*dd l} dt          j         }|                     |           d S t                      j        J d S )N)r,      r   r   zwYour version of hypothesis is outdated. To avoid `DeadlineExceeded` errors, please update. Current hypothesis version: )hypothesis_versionwarningsr3   __version__warnr   r   )r   warning_messages     r   assert_deadline_disabledr   e  sf    J&&D+5+AD D 	
 	o&&&&&zz"*****r   )NNNNN)r   Nr   NN)NNN)
r[   r   r   r   r   r   r   FNN).collectionsr   collections.abcr   r	   r   r   r3   	functoolsr   r   r   r   r:   hypothesis.extrarv   hypothesis.strategiesr
   (torch.testing._internal.common_quantizedr   r   quint8qint8qint32rF   rK   r)   r@   r;   rD   	compositer   rl   r   rC   r   r   r4   r5   r   	_profiles_current_profile__dict__current_settings	unlimitedr   r   r   r   <module>r      s   $ # # # # # $ $ $ $ $ $                            ' ' ' ' ' ' * * * * * * 0 0 0 0 0 0 w w w w w w w w 
L	K	L #{<<	L$	K	L!2   8 8 8& & &6, , ,
   9=04- - - -8 & & & &" tT & & & &0 , , , ,>.^ *0#)>E9>1 1 1 1l  '8 %h&?@I #  ##(:Y(F(F","6Y
+ 
+ 
+ 
+ 
+r   