
    Vhs              )       \   U d dl Z d dlm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mZ d dlmZmZ d dlmZmZmZmZmZmZmZ d d	lmZmZmZ d
Zer	 d dlZ dZ G d de      Z"d Z#d Z$	 d|dZ%d Z& edddejN                  jN                  ejP                   eejR                         eejR                  ge	sdnejT                  ejV                  f  ee%d      e#dddd
       eddejN                  jX                  dejZ                   eejR                         eejR                  ge	sdnejT                  ejV                  f  ee%d      e$dddd
 eej\                  dej^                  di      g e e j`                  d      ddd ejV                  ge!      f"       ed#d$dejN                  jb                  ejd                   eejR                         eejR                  ge	sdnejT                  ejV                  f  ee%d      e$dddd
 eej\                  dej^                  di      g%       ed&d'd(ejN                  jf                  ejP                   eejR                         eejR                  ge	sdnejT                  ejV                  f  ee%d)      e#dddd
d
 e e j`                  d      d*d+ejh                  ejj                  f,      f-       e	 d}d/d(ere jN                  jl                  ndejZ                   eejR                         eejR                  ge	sdnejT                  ejV                  f  ee%d0      e$dddd
d
 e eej\                  d1ej^                  d1i      d2d3      g e e j`                  d      d*d+       e e j`                  d      ddd 4      fd5 e	 d~d7d(ere jN                  jn                  ndejd                   eejR                         eejR                  ge	sdnejT                  ejV                  f  ee%d0      e$dddd
d
 e eej\                  d1ej^                  d1i      d2d3      g e e j`                  d      d*d+      fd5 ed8d9d:ejN                  jp                  ejP                   eejR                         eejR                  ge	sdnejT                  f  ee%d      e#dddd
dd
;       ed<d=d:ejN                  jr                  ejZ                   eejR                         eejR                  ge	sdnejT                  f  ee%d      e$dddd
d
 eej\                  di      g>       ed?d@d:ejN                  jt                  ejd                   eejR                         eejR                  ge	sdnejT                  f  ee%d      e$dddd
d
 eej\                  di      g>       edAdBdejN                  jv                  ejP                   ee%d      e#dddd
 eejR                         eejR                  ge	sdnejT                  ejV                  f C       edDdEdejN                  jx                  ejZ                   ee%d      e$dddd
 eejR                         eejR                  ge	sdnejT                  ejV                  f  e eej\                  dej^                  di      d2d3      gF       edGdHdejN                  jz                  ejd                   ee%d      e$dddd
 eejR                         eejR                  ge	sdnejT                  ejV                  f  e eej\                  dej^                  di      d2d3      gF       edIdJd:ejN                  j|                  ejP                   ee%d      e#ddd
 eejR                         eejR                  ge	sdnejT                  f dd
K       e	 ddMd:ere jN                  j~                  ndejZ                   ee%d      e$dddd
 eejR                         eejR                  ge	sdnejT                  f d
d
 ee j                  ddN       e eej\                  d1i      d2d3       ee j                  ddO       ee j                  ddP      fdQ e	 ddSd:ere jN                  j~                  ndejd                   ee%d      e$dddd
 eejR                         eejR                  ge	sdnejT                  f d
d
 ee j                  ddN       ee j                  ddO       e eej\                  d1i      d2d3      gdQ edTdUd(ejN                  j                  ejP                   ee%dV      e#dddd
 eejR                         eejR                  ge	sdnejT                  ejV                  f d
W       edXdYd(ejN                  j                  ejZ                   ee%dV      e$dddd
 eejR                         eejR                  ge	sdnejT                  ejV                  f d
 e eej\                  dej^                  di      d2d3      gZ       ed[d\d(ejN                  j                  ejd                   ee%dV      e$dddd
 eejR                         eejR                  ge	sdnejT                  ejV                  f d
 e eej\                  dej^                  di      d2d3      gZ       ed] eejR                  ej                  ejT                  ej                        e&d
dd^       ed_ eejR                  ej                  ejT                  ej                        e&d
dd^      gZFeGe   eHd`<    e"dadb       e"dcdAb       e"ddd8b       e"dedTb       e"dfd&b       e"dgdIb       e"dhd# e eej\                  dej^                  di      d2d3      gi       e"djdG e eej\                  dej^                  di      d2d3      gi       e"dkd?b       e"dld[ e eej\                  dej^                  di      d2d3      gi       e"dmd6 e eej\                  d1ej^                  d1i      d2d3      gi       e"dndR e eej\                  d1i      d2d3       e e j`                  d      ddoej                  fd p       e e j`                  d      ddqej                  fd p      gi       e"drdb       e"dsdD e eej\                  dej^                  di      d2d3      gi       e"dtd<b       e"dudX e eej\                  dej^                  di      d2d3      gi       e"dvd. e eej\                  d1ej^                  d1i      d2d3      gi       e"dwdL e eej\                  d1i      d2d3       e e j`                  d      ddqd 4      gi       edxeFd]y       edzeFd_y      gZJeGe   eHd{<   y# e!$ r Y w xY w)    N)partialmake_tensor)SM53OrLater)precisionOverride)all_types_andall_types_and_complex_and)
TEST_SCIPYTEST_WITH_ROCM)DecorateInfo
ErrorInputOpInfosample_inputs_spectral_opsSampleInputSpectralFuncInfoSpectralFuncType)_find_referenced_opinfo_inherit_constructor_argsPythonRefInfoFTc                   *     e Zd ZdZddd fd
Z xZS )SpectralFuncPythonRefInfozM
    An OpInfo for a Python reference of an elementwise unary operation.
    N )optorch_opinfo_variantc                    || _         t        ||t              | _        t	        | j                  t
              sJ | j                  j                  }t        ||||      }t        | $  di | y )N)op_db )
torch_opinfo_namer   r   torch_opinfo
isinstancer   _original_spectral_func_argsr   super__init__)	selfnamer   r   r   kwargs	inheritedukwargs	__class__s	           ^/home/dcms/DCMS/lib/python3.12/site-packages/torch/testing/_internal/opinfo/definitions/fft.pyr#   z"SpectralFuncPythonRefInfo.__init__0   sl     "3335
 $++-=>>>%%BB	+D"iH#7#    )__name__
__module____qualname____doc__r#   __classcell__)r)   s   @r*   r   r   +   s     $ $r+   r   c              +      K   t        t        |t        j                        }t	        t         |             t        d       y w)Ndevicedtypez6Dimension specified as -1 but tensor has no dimensions
error_typeerror_regexr   r   torchfloat32r   r   
IndexErrorop_infor3   r&   make_args       r*   error_inputs_fftr?   E   s5     {6GH
HJL s   AAc              +      K   t        t        |t        j                        }t	        t         |       d      t        d       y w)Nr2   )r   dimz5Dimension specified as 0 but tensor has no dimensionsr5   r8   r<   s       r*   error_inputs_fftnrC   O   s7     {6GH
HJD)K s   AAc             +      K   t        | |||fi |E d {    t        ry t        ||||      }t        |       y 7 (w)N)r4   r3   requires_grad)r   r   r   r   )r=   r3   r4   rE   min_sizer&   as          r*   sample_inputs_fft_with_minrH   Y   sT      *17    	HE&VA
a.s   A >)A c              +   >  K   fd}t         |d             t         |d      t        d             t         |d      t        d             t         |d	      t        d
             t         |d      t        d             y w)Nc                 $    t        | fd|S )N)r3   r4   rE   r   )shaper&   r3   r4   rE   s     r*   mtz"sample_inputs_fftshift.<locals>.mtj   s'    
 ]
NT
 	
r+   )	   
   )2   r   rA   )r&   )      )   )rP      )r   rR   )rP   rS      )r   rT   )r   dict)r=   r3   r4   rE   r&   rL   s    ```  r*   sample_inputs_fftshiftrV   i   ss     

 bk
""
bi
44
bk$4.
99
bj&)9
::
bmDV,<
==s   BBzfft.fftfft_fft_fft_c2cr   rR   )rF   )	aten_namedecomp_aten_namerefndimensionaldtypesdtypesIfCUDAsample_inputs_funcerror_inputs_funcgradcheck_fast_modesupports_forward_adsupports_fwgrad_bwgradcheck_batched_forward_gradzfft.fft2fft_fft2)rR   rR   g-C6?zSkipped!
TestCommon#test_complex_half_reference_testingcuda)device_typer]   	active_if)rY   r[   rZ   r\   r]   r^   r_   r`   ra   rb   rc   rd   
decoratorsskipszfft.fftnfft_fftn)rY   rZ   r[   r\   r]   r^   r_   r`   ra   rb   rc   rd   rk   zfft.hfftfft_hfft_fft_c2rrT   TestSchemaCheckModeOpInfotest_schema_correctness)r]   )rY   rZ   r[   r\   r]   r^   r_   r`   ra   rb   rc   rd   check_batched_gradgradrl   	fft.hfft2	fft_hfft2)rT   rT   g-C6*?TestFFTtest_reference_nd)ri   )rY   rZ   r[   r\   r]   r^   r_   r`   ra   rb   rc   rr   rd   rk   rl   	fft.hfftn	fft_hfftnzfft.rfftfft_rfft_fft_r2c)rY   rZ   r[   r\   r]   r^   r_   r`   ra   rb   rc   check_batched_gradrl   rr   z	fft.rfft2	fft_rfft2)rY   rZ   r[   r\   r]   r^   r_   r`   ra   rb   rc   r{   rr   rk   z	fft.rfftn	fft_rfftnzfft.ifftfft_ifft)rY   rZ   r[   r\   r_   r`   ra   rb   rc   rd   r]   r^   z	fft.ifft2	fft_ifft2)rY   rZ   r[   r\   r_   r`   ra   rb   rc   rd   r]   r^   rk   z	fft.ifftn	fft_ifftnz	fft.ihfft	fft_ihfft)rY   rZ   r[   r\   r_   r`   rb   rc   rd   r]   r^   rl   r{   
fft.ihfft2
fft_ihfft2test_out_warningtest_outtest_out_warnings)rY   rZ   r[   r\   r_   r`   ra   rb   rc   rd   r]   r^   r{   rr   rk   
fft.ihfftn
fft_ihfftnz	fft.irfft	fft_irfft)rR   rT   )rY   rZ   r[   r\   r_   r`   ra   rb   rc   rd   r]   r^   rr   z
fft.irfft2
fft_irfft2)rY   rZ   r[   r\   r_   r`   ra   rb   rc   rd   r]   r^   rr   rk   z
fft.irfftn
fft_irfftnzfft.fftshift)r]   r_   supports_outrb   rc   zfft.ifftshiftr   z_refs.fft.fft)r   z_refs.fft.ifftz_refs.fft.rfftz_refs.fft.irfftz_refs.fft.hfftz_refs.fft.ihfftz_refs.fft.fftn)r   rk   z_refs.fft.ifftnz_refs.fft.rfftnz_refs.fft.irfftnz_refs.fft.hfftnz_refs.fft.ihfftntest_python_ref_torch_fallback)r]   ri   test_python_ref_executorz_refs.fft.fft2z_refs.fft.ifft2z_refs.fft.rfft2z_refs.fft.irfft2z_refs.fft.hfft2z_refs.fft.ihfft2z_refs.fft.fftshift)r   r   z_refs.fft.ifftshiftpython_ref_db)F)rs   )rw   )r   )r   )Kunittest	functoolsr   numpynpr9   torch.testingr   #torch.testing._internal.common_cudar   *torch.testing._internal.common_device_typer   $torch.testing._internal.common_dtyper   r	   $torch.testing._internal.common_utilsr
   r   #torch.testing._internal.opinfo.corer   r   r   r   r   r   r   #torch.testing._internal.opinfo.refsr   r   r   has_scipy_fft	scipy.fftscipyModuleNotFoundErrorr   r?   rC   rH   rV   fftOneDboolhalf	complex32fft2TwoDfloatcfloatskipfftnNDhfft	complex64
complex128hfft2hfftnrfftrfft2rfftnifftifft2ifftnihfftihfftnexpectedFailureirfftirfft2irfftnbfloat16chalfr   list__annotations__float16r   r   r+   r*   <module>r      s       % ; H L    
$ 0 $4 +0 
> #FFJJ%**(4.JJ
$bEJJ+H
 ##=J*  ##('* FFKK#%**(4.JJ
$bEJJ+H
 ##=O+  ##(%u{{D%,,&MNOj)5"((	
+@ #FFKK%(((4.JJ
$bEJJ+H
 ##=O+  ##(%u{{D%,,&MNO), #FFKK%**(4.JJ
$bEJJ+H
 ##=J*  ##($ j)+))9)9:	
+> +#,EIIOO$%**(4.JJ
$bEJJ+H
 ##=O+  #$#(!5;;ellD"IJ#
 j)+) j)5"	
9+X $#,EIIOO$%(((4.JJ
$bEJJ+H
 ##=O+  #$#(!5;;ellD"IJ#
 j)+)
9$J #FFKK%**UZZ("JJ
%0"

}
 ##=J*  # $'* #FFLL%**UZZ("JJ
%0"

}
 ##=O+  # $u{{D12
'. #FFLL%((UZZ("JJ
%0"

}
 ##=O+  # $u{{D12
'. #FFKK%**"#=J*  ##((4.JJ
$bEJJ+H
!* #FFLL%**"#=O+  ##((4.JJ
$bEJJ+H

 !5;;ellD"IJ#
)8 #FFLL%(("#=O+  ##((4.JJ
$bEJJ+H

 !5;;ellD"IJ#
)8 #FFLL%**"#=O* ##(UZZ("JJ
%0"

}
  %( # -EII4%**"#=O+  ##(UZZ("JJ
%0"

}
 !$ 11<AST!5;;"56	CV 11<L11<ATU	
+@ # -EII4%(("#=O+  ##(UZZ("JJ
%0"

}
 !$ 11<AST11<L!5;;"56	CV
+> #FFLL%**"#=O*  ##((4.JJ
$bEJJ+H
  %), #FFMM%**"#=O+  ##((4.JJ
$bEJJ+H
  %!5;;ellD"IJ#
+: #FFMM%(("#=O+  ##((4.JJ
$bEJJ+H
  %!5;;ellD"IJ#
+: (JJ

EKK
 2 #	 (JJ

EKK
 2 #	sCtF| CL # $ $ % $ % $!5;;ellD"IJ#

 %!5;;ellD"IJ#

 % &!5;;ellD"IJ#

 %!5;;ellD"IJ#

 &!5;;"56# j)0'" j)*'"#
: $ %!5;;ellD"IJ#

 % &!5;;ellD"IJ#

 %!5;;ellD"IJ#

 &!5;;"56# j)*"	
( (
 )QmtF| mk  s   x" "x+*x+