
    bIg                    n    d dl mZ d dlZd dlZd dlmZ d dlmZ dej                  fdZ
 G d de      Zy)	    )annotationsN)Base)expectc                   t        j                  |       }t        |j                        }t        j                  |      }|dk  r||dz   z  }|j                  d| }|j                  || }t        j
                  |dt        |      z  |j                  z   dt        |      z  z         }	t        j
                  t        j                  ||      g |d|      }t        j                  |	|k(  |      S )z/Compute one hot from indices at a specific axisr      )r   dtype)npasarraylenshapearangereshapemod)
indicesdepthaxisr
   valuesrankdepth_rangelsrstargetss
             _/var/www/html/answerous/venv/lib/python3.12/site-packages/onnx/backend/test/case/node/onehot.pyone_hotr      s    ZZ Fv||D))E"Kaxq	a	B	d4	 BjjTCG^k&7&77$R.HG ZZvu-}}Q}}=F::g'u55    c                  T    e Zd Zedd       Zedd       Zedd       Zedd       Zy)OneHotc                 x   d} d}t         j                  }t        j                  j	                  dg ddg      }t        j
                  g dt         j                        }t        j                  d	      }t        j
                  || g|      }t        |||      }|| |z
  z  |z   }t        ||||g|gd
       y )N      r   r   r   r   y)inputsoutputs)r         r	      test_onehot_without_axisr%   r&   name)
r   int32onnxhelper	make_nodearrayint64float32r   r   )on_value	off_valueoutput_typenoder   r   r   r$   s           r   export_without_axiszOneHot.export_without_axis   s    	hh{{$$;cU % 
 ((9BHH5

29h/{CGU+6I%&2UF+C+		
r   c                    d} d}d}t         j                  }t        j                  j	                  dg ddg|       }t        j
                  ddgdd	ggt         j                  
      }t        j                  d      }t        j
                  ||g|
      }t        ||| |      }|||z
  z  |z   }t        ||||g|gd       y )Nr      r   r#   r$   r%   r&   r   	   r"      r	   
   r   r
   test_onehot_with_axisr+   r   r3   r.   r/   r0   r1   r   r   		axisValuer4   r5   r6   r7   r   r   r   r$   s	            r   export_with_axiszOneHot.export_with_axis2   s    		jj{{$$1E	 % 
 ((QFQF+2::>

29h/{CGU+FI%&2UF+C(		
r   c                    d} d}d}t         j                  }t        j                  j	                  dg ddg|       }t        j
                  g dt         j                        }t        j                  d	      }t        j
                  ||g|      }t        ||| |
      }|||z
  z  |z   }t        ||||g|gd       y )Nr   r:   r   r#   r$   r;   )r   iir	   r>   r?   test_onehot_negative_indicesr+   )	r   r3   r.   r/   r0   r1   r2   r   r   rB   s	            r   export_with_negative_indicesz#OneHot.export_with_negative_indicesJ   s    		jj{{$$1E	 % 
 ((;bhh7 

29h/{CGU+FI%&2UF+C/		
r   c                    d} d}d}t         j                  }t        j                  j	                  dg ddg|       }t        j
                  ddgd	d
ggt         j                        }t        j                  d      }t        j
                  ||g|      }t        ||| |      }|||z
  z  |z   }t        ||||g|gd       y )Nr:   r   r   r#   r$   r;   r<   r"   r=   r	   r>   r?   test_onehot_with_negative_axisr+   rA   rB   s	            r   export_with_negative_axisz OneHot.export_with_negative_axish   s    		jj{{$$1E	 % 
 ((QFQF+2::>

29h/{CGU+FI%&2UF+C1		
r   N)returnNone)__name__
__module____qualname__staticmethodr8   rD   rG   rK    r   r   r   r      sP    
 
& 
 
. 
 
: 
 
r   r   )
__future__r   numpyr   r.   onnx.backend.test.case.baser   onnx.backend.test.case.noder   r3   r   r   rR   r   r   <module>rW      s4    #   , . "$2:: 6 a
T a
r   