
    ΅i>:              	       @   S SK Jr  S SKJr  S SKrS SKrS SKrS SKJ	r	  S SK
Jr  S SKJr  S SKJr  S SKJr  S S	KJr  S S
KJr  S SKJrJr  \R.                  \R0                  \R2                  4r\" S \R.                  S\R0                  S\R2                  S 05      rS rS rS r S r \R@                    SS j5       r! \R@                  S S j5       r" \R@                  SSS\RF                  4S j5       r$\R@                  S!S j5       r% \R@                       S"S j5       r&\'" \(" \)\" S5      RU                  S5      SS 5      5      r+S\+s=::  a  S:  a  O  O\RX                  " S\RZ                  S9  O\RX                  " SSS9  \R\                  " S5        S r/g)#    )defaultdict)IterableN)reduce)version)assume)settings)
strategies)numpy)SearchStrategy)_calculate_dynamic_qparams&_calculate_dynamic_per_channel_qparamsc                      g )N r       b/home/james-whalen/.local/lib/python3.13/site-packages/torch/testing/_internal/hypothesis_utils.py<lambda>r   !   s    4r   c                    U u  pnS[         R                  " [         R                  5      R                  -   n[         R                  " [         R
                  5      nUR                  U-  UR                  U-  pv[        Xb-
  U-  Xa-  U-   5      n[        Xr-
  U-  Xq-  U-   5      n	[        R                  " U5      [        R                  " U	5      4$ 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11Jkk%**-O(,,z9?;N;NQ[;[hX*e3h6F6SUIX*e3h6F6SUI::i "**Y"777r   c                  $   SU;   Ga  [         R                  R                  S:  Ga  SU;   =(       a
    US   (       + =(       d    SU;  =(       a"    SU;   =(       a
    US   (       + =(       d    SU;  n[        U 5      S:H  =(       a    SU;  =(       a    SU;  nU(       GaK  U(       GaC  US   S:X  ac  [        R
                  " [        R                  5      R                  US'   [        R
                  " [        R                  5      R                  US'   OUS   S	:X  ac  [        R
                  " [        R                  5      R                  US'   [        R
                  " [        R                  5      R                  US'   OkUS   S
:X  ab  [        R
                  " [        R                  5      R                  US'   [        R
                  " [        R                  5      R                  US'   UR                  S5        [        R                  " U 0 UD6$ )Nwidth)   C   r   	allow_nanallow_infinityr   r(   r)          @   )
hypothesisr   __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@   4   s   &Z//@@:M V#?F;,?(? '&-&(I8H1I-I ,V+	 	 IN &v%&v% 	"
 7g"$&+kk%--&@&D&D{#&+kk%--&@&D&D{#B&&+kk%--&@&D&D{#&+kk%--&@&D&D{#B&&+kk%--&@&D&D{#&+kk%--&@&D&D{#

799d%f%%r   c                  .    SU;  a  SUS'   [        U 0 UD6$ )Nr,   r2   )r@   )r<   r=   s     r   r;   r;   O   s#    fwD+F++r   c                     [        U5      u  p#[        U R                  5       U:  5        [        U R                  5       U:*  5        g)NT)r*   r   r   r   )tensorr    r(   r)   s       r   assume_not_overflowingrD   c   s8    -g6I
6::<9$%
6::<9$%r   c                 <   Uc  [         n[        U[        [        45      (       d  U4nU " [        R
                  " U5      5      n[        R                  " U5      nUR                  UR                  p[        U   n
U
b  U
nO(Uc  UOUnUc  U	OUnU " [        R                  " XS95      nUc.  [        R                  " [        R                  5      R                  nUc.  [        R                  " [        R                  5      R                  nU " [        X#SS95      nXU4$ )N)r(   r)   r2   )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    x   s     ~ ftUm,,"//&12N^,J$ (7L!
(0$n(0$n"++KL
KK,00	KK,00	)KLEn,,r   c                   ^ US:  d   eUc  [        US-   S5      nUS:  d   eUc  US-   n[        R                  " [        R                  " X45      XS9nTb  UR	                  U4S j5      nU " UR                  [        5      5      $ )z8Return a strategy for array shapes (tuples of int >= 1).r2         )min_sizemax_sizec                 @   > [        [        R                  U S5      T:*  $ r   )r   int__mul__)x	max_numels    r   r   array_shapes.<locals>.<lambda>   s    vckk1a/HI/Ur   )r   r:   listsrL   filtermaprI   )rM   min_dimsmax_dimsmin_sidemax_siderc   	candidates        ` r   array_shapesrm      s     b==x!|R(b==a<X88_I$$%UV		e$%%r   c                 |   [        U[        5      (       a	  U " U5      nOU " [        R                  " U5      5      nUc  Uc  [	        SSSSS9nU " [
        R                  " XBUS95      n[        [        R                  " U5      R                  5       =(       d$    [        R                  " U5      R                  5       (       + 5        US 4$ U " U5      nUc  [        U5      u  px[	        XxSSSS9nU " [
        R                  " XBUS95      n[        XcS   5      u  p[        R                  US   S 5      nUb  Un
XiXS   44$ )	N    .    .AFr2   r/   r,   dtypeelementsshaper0   r/   r,   r[   )rG   r   r:   rJ   r;   stnparraysr   r   isnananyisinfr*   r   rK   get)rM   shapesrt   r    rs   _shapeXr(   r)   r!   zpenforced_zps               r   rC   rC      s   &.))fboof-.dC5CH56JKBHHQKOO%:!):;<$w7mG1':	)u$)5T[[uvFGA*1aj9IE&**71:t<Kb!*%%%r   c                    [        U[        5      (       a	  U " U5      nOU " [        R                  " U5      5      nUc  Uc  [	        SSSSS9nU " [
        R                  " [        R                  X$S95      n[        [        R                  " U5      R                  5       =(       d$    [        R                  " U5      R                  5       (       + 5        US 4$ U " U5      nUc  [        U5      u  pg[	        XgSSSS9nU " [
        R                  " [        R                  X$S95      n[        XSS   5      u  p[        R!                  US   S 5      n
U
b  U
n	[#        [        R$                  R'                  S	UR(                  S
5      5      n[        R*                  " UR(                  5      nXS	'   S	X'   [        R,                  " X\5      nXXXUS   44$ )Nro   rp   Fr2   rq   rr   rv   r[   r   r   )rG   r   r:   rJ   r;   rw   rx   r   r   r   ry   rz   r{   r*   r   rK   r|   r`   randomrandintndimarange	transpose)rM   r}   rt   r    r~   r   r(   r)   r!   r   r   axispermute_axess                r   per_channel_tensorr      ss   &.))fboof-.dC5CH2::OPBHHQKOO%:!):;<$w7mG1':	)u$)5T[[rzzHKLA6q!*EIE&**71:t<Kryy  AFFA./D99QVV$LOL
Q%Ab
+++r   c                    U " [         R                  " U6 5      nU " [         R                  " U6 5      nU " [         R                  " U6 5      nU " [         R                  " SU5      5      nX-  nX-  n[        U[        5      (       a  U " [         R                  " U5      5      n[        U5       Vs/ s H  nU " [         R                  " U6 5      PM     nn[        U5       Vs/ s H  nU " [         R                  " U6 5      PM     nnSnUU4[        U5      -   nUnU(       a3  U " [         R                  " 5       5      nU(       a  X4[        U5      -   nUnU
b8  [        U
[        [        45      (       a  [        U
5      S:X  d   S5       eOU
/S-  n
U " [        X4[        U5      -   4XS   S95      nU " [        U4U	U
S   S95      nU " [        U4U	U
S   S95      nUUUUU4$ s  snf s  snf )Nr   Fr-   zNeed 3 qparams for X, w, br   )r}   rt   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_rangefeature_map_rangekernel_range
max_groupscan_be_transposedrt   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   *  s    bkk#345J#
34 6 $
45!7"++a,-F-6N/8O+x((2??;78HMkHZ[HZ1bkk+<=>HZ[9>{9KL9KAtBKK./9KGL	B#%=>wOL J"++- *FwWL(J ge}--w<1$B&BB$i!mGV	$u->'??B1:/ 	0A 	VL?X#AJ( 	)AVJ=8#AJ( 	)A aFB9 \Ls   /#G!!#G&r4   .r-   )r-   r1   r   r-      r   no_deadline)timeout)deadlinec                      [         S:  a(  SSKn S[        R                   3nU R	                  USS9  g[        5       R                  b   eg)zICheck that deadlines are effectively disabled across Hypothesis versions.r   r   NzwYour version of hypothesis is outdated. To avoid `DeadlineExceeded` errors, please update. Current hypothesis version: r[   )
stacklevel)hypothesis_versionwarningsr4   __version__warnr   r   )r   warning_messages     r   assert_deadline_disabledr   o  sO    J&++5+A+A*BD 	
 	o!4z""***r   )NNNNN)r   Nr   NN)NNN)
r[   )r      r-      r   )      r   r   FNN)0collectionsr   collections.abcr   r
   r   r   r4   	functoolsr   importlib.metadatar   r   r   r	   r:   hypothesis.extrarw   hypothesis.strategiesr   (torch.testing._internal.common_quantizedr   r   quint8qint8qint32rF   rK   r*   r@   r;   rD   	compositer    rm   r   rC   r   r   rI   rg   r`   splitr   register_profile	unlimitedload_profiler   r   r   r   <module>r      s   $ $     &   ' * 0 w 
LL	KK	LL #<	LL$	KK	LL!2  8&6,
 9=04- -8 & &" tT & &0 , ,>.^ *0#)>E9>1 1n 3sGL$9$?$?$DRa$HIJ #0j0mZ5I5IJ md; 	  m $+r   