
    bIg                    v    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
mZ d Zd Zd Zd	 Z G d
 de      Zy)    )annotationsN)Base)expect)apply_affine_transformconstruct_original_gridc                   t        j                  t        j                  |       t        j                  |       t        j                  |       t        j                  |       t        j                  |       t        j
                  |        t        j                  |       t        j
                  |       t        j                  |       g	d      j                  ddd      }t        j                  t        j                  |      t        j                  |      t        j
                  |      t        j                  |      t        j                  |      t        j                  |      t        j
                  |       t        j                  |      t        j                  |      g	d      j                  ddd      }t        j                  t        j                  |      |||t        j                  |      |||t        j                  |      g	d      j                  ddd      }t        j                  |t        j                  |      t        j                  |      t        j                  |      |	t        j                  |      t        j                  |      t        j                  |      |
g	d      j                  ddd      }t        j                  t        j                  |||g            j                  ddd      }||z  |z  |z  }t        j                  |d      }t        j                  ||f      }t        j                  |d      }|j                  t         j                        S )Naxis      r      r   )npstack	ones_like
zeros_likecossinreshape	transposearrayhstackastypefloat32)angle1angle2offset_xoffset_yoffset_zshear_xshear_yshear_zscale_xscale_yscale_zrot_xrot_yshearscaletranslationrotation_matrixaffine_matrixs                     c/var/www/html/answerous/venv/lib/python3.12/site-packages/onnx/backend/test/case/node/affinegrid.pycreate_affine_matrix_3dr/      s    HHLL MM&!MM&!MM&!FF6NVVF^OMM&!FF6NFF6N
	
  gb!Q 
 HHFF6NMM&!FF6NMM&!LL MM&!VVF^OMM&!FF6N
	
  gb!Q 
 HHLL!LL!LL!
	
  gb!Q 
 HHMM'"MM'"MM'"MM'"MM'"MM'"
	
  gb!Q 
 ,,rxx8X(FGHPP
AqK eme+e3Oll?I>OII<=MLL	:M

++    c                   t        j                  t        j                  |       t        j                  |        t        j                  |       t        j                  |       gd      j	                  ddd      }t        j                  t        j
                  |      ||t        j
                  |      gd      j	                  ddd      }t        j                  |t        j                  |      t        j                  |      |gd      j	                  ddd      }	t        j                  t        j                  ||g            j	                  ddd      }
||z  |	z  }t        j                  |d      }t        j                  ||
f      }t        j                  |d      }|j                  t         j                        S )Nr	   r
   r   r   r   )r   r   r   r   r   r   r   r   r   r   r   r   )r   r   r   r!   r"   r$   r%   rotr)   r*   r+   r,   r-   s                r.   create_affine_matrix_2dr3   `   sZ    ((	"&&."&&."&&.IPRgb!Q  HH	g	",,w2GHrgb!Q 
 HH	"--("--*@'JQSgb!Q 
 ,,rxx8(<=>FFr1aPKEkE)Oll?I>OII<=MLL	:M

++r0   c            	        t        j                  t         j                  dz  t         j                  dz  g      } t        j                  ddg      }t        j                  ddg      }t        j                  ddg      }t        j                  d	d
g      }t        j                  ddg      }t        j                  ddg      }t        | ||||||      }|S )N   r         @      @ffffff
皙?            ?333333?333333ӿ皙@@?)r   r   pir3   )angler   r   r!   r"   r$   r%   theta_2ds           r.   create_theta_2drD   t   s    HHbeeai+,Exxc
#Hxxs$Hhhc{#GhhT{#GhhSz"GhhSz"G&x7GWgH Or0   c                    t        j                  t         j                  dz  t         j                  dz  g      } t        j                  t         j                  dz  t         j                  dz  g      }t        j                  ddg      }t        j                  ddg      }t        j                  d	d
g      }t        j                  ddg      }t        j                  ddg      }t        j                  ddg      }t        j                  d
dg      }t        j                  ddg      }	t        j                  ddg      }
t        | |||||||||	|
      }|S )Nr5   r      r   r6   r7   r8   r9   gr>   r:   r;   r<   r=   gffffff?gɿr?   r@   g      ?)r   r   rA   r/   )r   r   r   r   r    r!   r"   r#   r$   r%   r&   theta_3ds               r.   create_theta_3drH      s   XXruuqy"%%!),-FXXruuqy"%%!),-Fxxc
#Hxxs$Hxxs$Hhhc{#GhhT{#GhhT{#GhhSz"GhhSz"GhhSz"G&H Or0   c                  0    e Zd Zedd       Zedd       Zy)
AffineGridc                 \   t               } t        |       dddf\  }}}}||f}dD ]  }t        j                  j	                  dddgdg|	      }t        ||      }t        | |      }	d
}
|dk(  r|
dz  }
t        || t        j                  ||||gt        j                        g|	g|
        y )Nr      rF   r   r   rJ   thetasizegridinputsoutputsalign_cornerstest_affine_grid_2dr   _align_cornersdtyperR   rS   name)rD   lenonnxhelper	make_noder   r   r   r   r   int64)rC   NCHW	data_sizerT   nodeoriginal_gridrP   	test_names              r.    export_2d_no_reference_evaluatorz+AffineGrid.export_2d_no_reference_evaluator   s    "$]Aq!+
1aF	# 	M;;(((+	 ) D 4I}MM)(MBD-I!--	 "((Aq!Q<rxx"HI		r0   c                 d   t               } t        |       ddddf\  }}}}}|||f}dD ]  }t        j                  j	                  dddgd	g|
      }t        ||      }	t        | |	      }
d}|dk(  r|dz  }t        || t        j                  |||||gt        j                        g|
g|        y )Nr   r5   rL   rF   rM   rJ   rN   rO   rP   rQ   test_affine_grid_3dr   rV   rW   rY   )rH   r[   r\   r]   r^   r   r   r   r   r   r_   )rG   r`   ra   Drb   rc   rd   rT   re   rf   rP   rg   s               r.    export_3d_no_reference_evaluatorz+AffineGrid.export_3d_no_reference_evaluator   s    "$Hq!Q11aA1I	# 	M;;(((+	 ) D 4I}MM)(MBD-I!--	 "((Aq!Q?"(("KL		r0   N)returnNone)__name__
__module____qualname__staticmethodrh   rl    r0   r.   rJ   rJ      s(     2  r0   rJ   )
__future__r   numpyr   r\   onnx.backend.test.case.baser   onnx.backend.test.case.noder   !onnx.reference.ops.op_affine_gridr   r   r/   r3   rD   rH   rJ   rs   r0   r.   <module>ry      s=    #   , .L,^,(:3 3r0   