
    sh                        S SK r S SKrS SKrS SKrS SKJrJr  S SKJ	r	J
r
JrJrJr  S SKJrJrJrJrJr  S SKJr  S SKJrJr  S SKJr  S SKJr  S S	KJr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5      r+S r, " S S\5      r- " S S\R\                  \R^                  \5      r0\!Rb                  Re                  \05        g)    N)jittypeof)cgutilstypes	serializesigutilserrors)	is_jittedoverload_attributeoverload_methodregister_jitable	intrinsic)npydecl)AbstractTemplate	signature)tuple_setitem)	_internal)	UfuncBaseUfuncLowererBase)array_analysis)ufuncbuilder)numpy_support)Callable)ir)global_compiler_lockc                   N    \ rS rSrSS jrS rS rS rS rS r	S	 r
S
 rS rSrg)UfuncAtIterator   Nc                 X    Xl         X l        X0l        X@l        XPl        X`l        Xpl        g N)ufuncaa_tyindices
indices_tybb_ty)selfr!   r"   r#   r$   r%   r&   r'   s           i/Users/tiagomarins/Projetos/claudeai/copy_bank/venv/lib/python3.13/site-packages/numba/np/ufunc/dufunc.py__init__UfuncAtIterator.__init__   s%    
	$	    c                     U R                  X5        U R                  R                  5       u  p4U R                  XU5        U R                  R	                  5         g r    )_prepareindexerbegin_loops_call_ufunc	end_loops)r(   contextbuilderloop_indices_s        r)   runUfuncAtIterator.run$   sB    g',,224<8 r,   c                 J    [        U R                  [        R                  5      $ r    )
isinstancer%   r   	BaseTupler(   s    r)   need_advanced_indexing&UfuncAtIterator.need_advanced_indexing*   s    $//5??;;r,   c                 @   SSK JnJn  U R                  U R                  peU R
                  U R                  pUR                  [        R                  5      " S5      n	U R                  b  [        R                  " X)5      U l        U R                  5       (       a6  [        R                  " X&[!        U5      S9nUR                  n
U" XX5      u  pOU4nU4n
U" XX5      u  pU" XXuX5      U l        U R"                  R%                  5         U R'                  X5      U l        g )Nr   )normalize_indicesFancyIndexer)count)numba.np.arrayobjr@   rA   r"   r$   r#   r%   get_value_typer   intpr&   r   alloca_once_valueb_indicer=   unpack_tuplelenr/   prepare_compile_ufunccres)r(   r3   r4   r@   rA   r"   r$   r#   r%   zeroindex_typess              r)   r.   UfuncAtIterator._prepare-   s    EVVT\\799dooj%%ejj1!466#55gDDM&&((**714ZBG$**K#4W5@$K K jG%-K#4W5@$K K $Gd$/:''9	r,   c           	         SSK Jn  [        R                  " X$R                  5      n[        R                  " X$R
                  5      nUR                  n	[        R                  " XXUUR                  U5      n
U" XXZ5      nX4$ )Nr   )	load_item)	rC   rQ   r   rH   shapestridesdataget_item_pointer2layout)r(   r3   r4   r5   arrayarray_tyrQ   shapesrS   rT   ptrvals               r)   	_load_valUfuncAtIterator._load_valI   si    /%%g{{;&&w>zz''$(0G(8xr,   c                 N   UR                  U5      nUR                  U[        R                  5      nUR	                  [
        R                  U5      nU" X$U45      n	UR                  [        R                  5      " S5      n
UR                  Xj5      nUR                  Xc5        S U	4$ )N   )
loaddtyper   rE   get_functionoperatorgetitemrD   addstore)r(   r3   r4   r$   rW   rX   idxsigimplr[   ones              r)   
_load_flatUfuncAtIterator._load_flatT   s    ll7#nnXuzz2##H$4$4c:7CL) $$UZZ03kk##c#Syr,   c                     SSK Jn  U R                  R                  R                  nUR
                  n	UR                  X&X5      nU" XXFU5        g )Nr   )
store_item)rC   rn   rL   r   return_typera   cast)
r(   r3   r4   rW   rX   rZ   r[   rn   fromtytotys
             r)   
_store_valUfuncAtIterator._store_vala   s?    0$$00~~ll767XC8r,   c                 4   U R                   R                  S   nU R                  c  U R                  R                  4nO,U R                  R                  U R
                  R                  4nUR                  U5      nUR                  UR                  45        U$ Nr   )	r!   keyr&   r#   ra   r'   re   add_linking_libslibrary)r(   r3   r4   r!   rh   rL   s         r)   rK   UfuncAtIterator._compile_ufunch   sm    

q!66>99??$C99??DIIOO4Cyy~  $,,1r,   c                    U R                   nU R                  U R                  peU R                  XX5U5      u  pxU R                  c  U4n	O;U R                  U R
                  U R                  pn
U R                  XXU5      u  pX4n	UR                  X$R                  UR                  U	5      nU R                  XXVX5        g r    )rL   r"   r#   r\   r&   r'   rG   rk   call_internalfndescr   rs   )r(   r3   r4   r5   rL   r"   r#   rZ   r[   argsr&   r'   b_idxr6   val_bress                   r)   r1   UfuncAtIterator._call_ufunct   s    yy&&$))4>>'LTJ66>6D!VVTYYUAw4HHA<D##G[[$..$(*!3<r,   )
r"   r#   r&   rG   r'   rL   r/   r$   r%   r!   NN)__name__
__module____qualname____firstlineno__r*   r7   r=   r.   r\   rk   rs   rK   r1   __static_attributes__ r,   r)   r   r      s/    !<:8	9
=r,   r   c                    ^  SSK Jn   " U 4S jSUR                  5      nU=R                  T R                  R                  -  sl        U$ )Nr   npyimplc                   4   >^  \ rS rSrSr YrU 4S jrSrU =r$ )(make_dufunc_kernel.<locals>.DUFuncKernel   z
npyimpl._Kernel subclass responsible for lowering a DUFunc kernel
(element-wise function) inside a broadcast loop (which is
generated by npyimpl.numpy_ufunc_kernel()).
c                    > [         TU ]  XU5        U R                  R                  UR                  5      u  U l        U l        g r    )superr*   dufuncfind_ewise_functionr~   	inner_sigrL   )r(   r3   r4   	outer_sig	__class__s       r)   r*   1make_dufunc_kernel.<locals>.DUFuncKernel.__init__   s6    GWy9(,(G(G) %DNDIr,   )rL   r   )	r   r   r   r   __doc__r   r*   r   __classcell__)r   _dufuncs   @r)   DUFuncKernelr      s    	
 	  	 r,   r   )numba.npr   _Kernelr   r!   )r   r   r   s   `  r)   make_dufunc_kernelr      s9      w   W]]333r,   c                   ,   ^  \ rS rSrSrU 4S jrSrU =r$ )DUFuncLowerer   zHCallable class responsible for lowering calls to a specific DUFunc.
    c                 P   > SSK Jn  [        TU ]  U[        UR
                  5        g )Nr   r   )r   r   r   r*   r   numpy_ufunc_kernel)r(   r   r   r   s      r)   r*   DUFuncLowerer.__init__   s!    $+ 33	5r,   r   )r   r   r   r   r   r*   r   r   r   s   @r)   r   r      s    5 5r,   r   c                   n  ^  \ rS rSrSr\" S5      rSS0 4S jrU 4S jrS r	\
S	 5       rS
 r\S 5       r\S 5       r\S 5       r\S 5       r\S 5       r\S 5       r\S 5       r\S 5       rS rS rU 4S jrS r\S"S j5       rS rS#S jrS#S jrS#S jrS#S jr S"U 4S jjr!S"S jr"S r#S  r$S!r%U =r&$ )$DUFunc   z
Dynamic universal function (DUFunc) intended to act like a normal
Numpy ufunc, but capable of call-time (just-in-time) compilation
of fast loops specialized to inputs.
)identity
_keepaliveninnoutNFc                    [        U5      (       a  UR                  n[        R                  " 5          [	        SSUS.UD6" U5      nS S S 5        U R                  WU5        [        R                  " X5        g ! , (       d  f       N7= f)Nnpyufunc)_targetcacher   )r
   py_funcr   3_suppress_deprecation_warning_nopython_not_suppliedr   _initialize	functoolsupdate_wrapper)r(   r   r   r   targetoptions
dispatchers         r)   r*   DUFunc.__init__   sv    WooGMMO .Z#(.,..57J P 	X.  / POs   A88
Bc                 P  > [         R                  " U5      n[        [        U ]  XS9  U[
        R                  :g  U l        UR                  R                  U l	        UR                  R                  U l
        [        U 5      U l        U R                  5         U R                  5         g )N)r   )r   parse_identityr   r   r*   r   PyUFunc_Nonereorderabler   r   r   r   	_lower_me_install_cg_install_type)r(   r   r   r   s      r)   r   DUFunc._initialize   s    ..x8fd$Z$C %	(>(>>"**33!))11&t,r,   c                     [        U R                  R                  R                  5       5      n[	        U R                  U R
                  U R                  US9$ )$
NOTE: part of ReduceMixin protocol
)r   r   frozensiglist)list_dispatcher	overloadskeysdictr   _frozen)r(   r   s     r)   _reduce_statesDUFunc._reduce_states   sJ     t''116689'']]<<	
 	
r,   c                     [         R                  R                  U 5      nUR                  X5        U H  nUR	                  U5        M     U(       a  UR                  5         U$ )r   )r   _DUFunc__new__r   re   disable_compile)clsr   r   r   r   r(   rh   s          r)   _rebuildDUFunc._rebuild   sQ    
   ((-.CHHSM   "r,   c                     U $ )z;
For compatibility with the various *UFuncBuilder classes.
r   r<   s    r)   build_ufuncDUFunc.build_ufunc   s	     r,   c                 .    U R                   R                  $ r    )r   r   r<   s    r)   r   DUFunc.targetoptions   s    ---r,   c                 .    U R                   R                  $ r    )r!   r   r<   s    r)   r   
DUFunc.nin   s    zz~~r,   c                 .    U R                   R                  $ r    )r!   r   r<   s    r)   r   DUFunc.nout   s    zzr,   c                 .    U R                   R                  $ r    )r!   nargsr<   s    r)   r   DUFunc.nargs       zzr,   c                 .    U R                   R                  $ r    )r!   ntypesr<   s    r)   r   DUFunc.ntypes   s    zz   r,   c                 .    U R                   R                  $ r    )r!   r   r<   s    r)   r   DUFunc.types   r   r,   c                 .    U R                   R                  $ r    )r!   r   r<   s    r)   r   DUFunc.identity   s    zz"""r,   c                 .    U R                   R                  $ r    )r!   r   r<   s    r)   r   DUFunc.signature   s    zz###r,   c                 \    [        U R                  R                  5      S:  d   eSU l        g)z9
Disable the compilation of new signatures at call time.
r   TN)rI   r   r   r   r<   s    r)   r   DUFunc.disable_compile  s*    
 4##--.222r,   c                 T    [         R                  " U5      u  p#U R                  X#5      $ )z-
Compile the DUFunc for the given signature.
)r   normalize_signature_compile_for_argtys)r(   rh   r~   ro   s       r)   re   
DUFunc.add  s(     %88=''::r,   c                 0  > [         R                  R                  R                  nU[	        UR                  5       5      -    HB  n[        [        U5      SU5      ULd  M  UR                  " U S/UQ70 UD6nU[        Ld  M@  Us  $    [        TU ](  " U0 UD6$ )ze
Allow any argument that has overridden __array_ufunc__ (NEP-18)
to take control of DUFunc.__call__.
__array_ufunc____call__)r   npndarrayr   tuplevaluesgetattrtypeNotImplementedr   r   )r(   r~   kwsdefaultargoutputr   s         r)   r   DUFunc.__call__  s    
  ""**::%

--CtCy"3W=WL,,T:LLL/!M	 . 7#T1S11r,   c           	      V   U R                   R                  nU(       aP  SU;   a  UR                  S5      nX4-  nU(       a-  [        SSR	                  S [        U5       5       5      -  5      e[        U5      nXS:X  d  XSU R                   R                  -   :X  d   eU(       a   e/ nUS U  Ha  n[        U5      n[        U[        R                  5      (       a  UR                  nO[        R                  " U5      nUR                  U5        Mc     U R!                  [#        U5      5      $ )Noutz)unexpected keyword arguments to ufunc: %sz, c              3   8   #    U  H  n[        U5      v   M     g 7fr    )repr).0ks     r)   	<genexpr>+DUFunc._compile_for_args.<locals>.<genexpr>*  s     +I[DGG[   )r!   r   pop	TypeErrorjoinsortedrI   r   r   r:   r   Arrayra   r   map_arrayscalar_typeappendr   r   )	r(   r~   r   r   r   args_lenargtysr   argtys	            r)   _compile_for_argsDUFunc._compile_for_args"  s    jjnn|ggen K"&))+IVC[+I"I!J K K t9Xtzz1F%FGGw:C3KE%-- &::3?MM%   ''f66r,   c                     U R                   (       a  [        SU < 35      e[        U[        5      (       d   eUc  UnOU" U6 nU R                  R
                  R                  5        HC  u  pEXR                  :X  d  M  SU S3n[        R                  " U[        R                  5        Us  $    [        R                  " U R                  U R                  U5      u  pQn[        R                  " XQU5      n[        R                   " XW5      u  pn
U R#                  [%        U	5      U5        U R&                  R)                  XR*                  U
45        U R,                  R.                  R)                  UR*                  5        U$ )z
Given a tuple of argument types (these should be the array
dtypes, and not the array types themselves), compile the
element-wise function for those inputs, generate a UFunc loop
wrapper, and register the loop with the Numpy ufunc object for
this DUFunc.
compilation disabled for z>Compilation requested for previously compiled argument types (z~). This has no effect and perhaps indicates a bug in the calling code (compiling a ufunc more than once for the same signature)r   RuntimeErrorr:   r   r   r   itemsr~   warningswarnr	   NumbaWarningr   _compile_element_wise_functionr   _finalize_ufunc_signature!_build_element_wise_ufunc_wrapper	_add_loopintr   r  ry   r   libs)r(   r  ro   rh   r  rL   msg
actual_sig	dtypenumsrZ   envs              r)   r   DUFunc._compile_for_argtysA  sD    <<FGG&%((((Cv&C''11779GA""( *EE c6#6#67 : %1$O$Od00#%7!k!;;+'
*LL	s3x+\\378""4<<0r,   c                      UR                   U:H  $ r    )r~   )r(   ewise_typesrh   s      r)   match_signatureDUFunc.match_signaturef  s    xx;&&r,   c                     S[         S[        4S jn[        R                  " U5      nSnU H  nU" U5      n[	        X55      " U5        M     g )Nattrreturnc                    ^  U 4S jnU$ )Nc                 D   >^ [        U R                  S   T5      mU4S j$ )Nr   c                    > T$ r    r   )r!   r[   s    r)   <lambda>UDUFunc._install_ufunc_attributes.<locals>.get_attr_fn.<locals>.impl.<locals>.<lambda>o  s    Sr,   )r   rw   )r!   r[   r+  s    @r)   ri   CDUFunc._install_ufunc_attributes.<locals>.get_attr_fn.<locals>.implm  s    eiilD1((r,   r   )r+  ri   s   ` r)   get_attr_fn5DUFunc._install_ufunc_attributes.<locals>.get_attr_fnk  s    ) Kr,   )r   r   r   r   r   )strr   r   Functionr   )r(   templater3  at
attributesr+  attr_fns          r)   _install_ufunc_attributes DUFunc._install_ufunc_attributesi  sM    	c 	h 	 ^^H%/
D!$'Gr(1 r,   c                 H    U R                  U5        U R                  U5        g r    )_install_ufunc_reduce_install_ufunc_at)r(   r7  s     r)   _install_ufunc_methodsDUFunc._install_ufunc_methods~  s    ""8,x(r,   c                 `   ^  [         R                  " U5      n[        US5      SU 4S jj5       ng )Nr8  c                   >^ ^^^^ [         R                  " S[        R                  S9  [	        U[
        R                  5      (       d  Sn[        R                  " U5      e[	        U[
        R                  5      n[	        U[
        R                  5      n[	        U[
        R                  5      n[	        U[
        R                  5      nU=(       d    U=(       d    U(       + n	U=(       a    [        U5      S:H  n
[	        U[
        R                  [
        R                  [
        R                  [
        R                  45      n[        R                  " U5      mU=(       d    T(       + n[        U/5      nTR                   R"                  nUS:X  a1  [        R                  " U5      (       a  [        R$                  " S5      eUS:X  a3  [        R                  " U5      (       d  Sn[        R$                  " U5      e[        R                  " U5      (       a  TR'                  UR(                  45        OLU(       a  TR'                  UR(                  U45        O'TR'                  UR(                  UR(                  45        U 4S	 jm[*        U4S
 j5       m[*        U4S j5       mSU4S jjnSUUU4S jjnSS jnSU4S jjnU(       a  U$ U
(       a	  U(       a  U$ U	(       a	  U(       a  U$ U$ )Nz ufunc.at feature is experimentalcategoryz)The first argument "a" must be array-liker      zsecond operand needed for ufuncr_   z+second operand provided when ufunc is unaryc           	         > SSK Jn  [        U5      S:X  a  UR                  u  pVpxUu  pYpOUR                  S-   u  pVpxUS-   u  pYpU" U5      " XU	5      n	[	        TXXX5      nUR                  X5        g )Nr   )
make_array   r    )rC   rH  rI   r~   r   r7   )r3   r4   rh   r~   rH  r6   atyidxtybtyr"   r$   r&   at_iterr!   s                r)   apply_ufunc_codegenDDUFunc._install_ufunc_at.<locals>.ol_at.<locals>.apply_ufunc_codegen  sx    8t9>),&AE'+$A'1),G);&AE'+g~$A'sOGa8)%PG-r,   c                 :   > [         R                  " XX45      nUT4$ r    r   none)	typingctxr!   r"   r$   r&   rh   rN  s         r)   apply_a_b_ufunc@DUFunc._install_ufunc_at.<locals>.ol_at.<locals>.apply_a_b_ufunc  s    jj76///r,   c                 :   > [         R                  " XU5      nUT4$ r    rQ  )rS  r!   r"   r$   rh   rN  s        r)   apply_a_ufunc>DUFunc._install_ufunc_at.<locals>.ol_at.<locals>.apply_a_ufunc  s    jj73///r,   c                    > T(       a&  U R                  U[        R                  " U5      5      $ U R                  U[        R                  " U5      [        R                  " U5      5      $ r    )r8  r   asarrayr!   r"   r$   r&   b_nones       r)   	impl_cast:DUFunc._install_ufunc_at.<locals>.ol_at.<locals>.impl_cast  sJ     88Arzz'':;; 88A$&JJw$7$&JJqM3 3r,   c                    > T(       a
  T" XU5        g [         R                  " U5      nX   n[         R                  " XER                  5      nT" XX$R                  5        g r    )r   rZ  broadcast_torR   flat)	r!   r"   r$   r&   b_a_rT  rW  r\  s	         r)   impl_generic=DUFunc._install_ufunc_at.<locals>.ol_at.<locals>.impl_generic  sD    !%G5ABBXX6B#Egww?r,   c                 "    U " US   U5      US'   g )Nr   r   )r!   r"   r$   r&   s       r)   impl_indices_empty_b_scalarLDUFunc._install_ufunc_at.<locals>.ol_at.<locals>.impl_indices_empty_b_scalar  s    aeQ"r,   c                 H   > T(       a  U " X   5      X'   g U " X   U5      X'   g r    r   r[  s       r)   impl_scalar_scalarCDUFunc._install_ufunc_at.<locals>.ol_at.<locals>.impl_scalar_scalar  s$    !&qz!2AJ!&qz1!5AJr,   r    )r  r  r	   NumbaExperimentalFeatureWarningr:   r   r  NumbaTypeErrorListTuple	SliceTyperI   Sequencer   is_nonelikeanyr!   r   TypingErrorre   ra   r   )r!   r"   r$   r&   r!  indices_arrindices_listindices_tupleindices_sliceindices_scalarindices_empty_tupleb_arrayb_scalar	need_castr   r]  rd  rg  rj  rT  rW  rN  r\  r(   s   `                  @@@@r)   ol_at'DUFunc._install_ufunc_at.<locals>.ol_at  s'   MM<#)#I#IK a--A++C00$Wekk:K%guzz:L&w<M&w@M"-"O"O-PN"/"ECLA4E U[[%..%**%*[[%2 3G((+F#-v.H\N+I**..C axG//22(()JKK ax 3 3A 6 6C((--""1%%!''$!''1&!''177+,. 0 0 0 03@ @(6   $22H))##r,   r    r   r6  r   )r(   r7  r8  r~  s   `   r)   r?  DUFunc._install_ufunc_at  s,    ^^H%	T	"_	$ 
#_	$r,   c                 `   ^  [         R                  " U5      n[        US5      SU 4S jj5       ng )Nreducec                    >^^^^^^^^^^^^^^ [         R                  " S[        R                  S9  [	        U[
        R                  5      (       d  Sn[        R                  " U5      e[	        U[
        R                  5      =(       a$    [	        UR                  [
        R                  5      n[	        U[
        R                  5      =(       a    [        U5      S:H  n[        R                  " U5      nTR                  R                   S L mTR                  R"                  mTR$                  (       d2  U(       a+  [        U5      S:  a  ST S3n[        R                  " U5      eSUR&                  -  mSUR&                  S-
  -  m[        R                  " U5      (       a  UR                  OUmTR                   m[        R                  " T5      m[        R                  " U5      m[(        U4S	 j5       m[(        U4S
 j5       m[*        S 5       m[(        S 5       m[(        S 5       mSUUUUU4S jjn	   SUUUUUUU4S jjn
   SUUU4S jjn   SS jn   SU4S jjnUR&                  S:X  a	  U(       d  U	$ U(       a  U$ U(       a   [-        [/        UR&                  5      5      mU$ U(       a  S[        U5      S-
  -  mU$ US:X  d>  [	        U[
        R                  [
        R0                  [
        R2                  45      (       a  U
$ g )Nz$ufunc.reduce feature is experimentalrD  z-The first argument "array" must be array-liker   r_   zreduction operation 'z:' is not reorderable, so at most one axis may be specified)r   c                 h   > TnSn[        U 5       H  u  pEXA:X  a  M  [        X#U5      nUS-  nM     U$ Nr   r_   )	enumerater   )tuppossijetup_init_m1s         r)   tuple_sliceDDUFunc._install_ufunc_reduce.<locals>.ol_reduce.<locals>.tuple_slice  sC      %cNDAx %aA.AFA	 +
 r,   c                    > TnSS[        U5      penXV:  a3  XQ:X  a  [        X5U5      nOX   n[        X5U5      nUS-  nUS-  nXV:  a  M3  U$ r  )rI   r   )	r  r  r[   r  r  r  szr  tup_inits	           r)   tuple_slice_appendKDUFunc._install_ufunc_reduce.<locals>.ol_reduce.<locals>.tuple_slice_append  sb     aQbfx)!4F)!2QFA f r,   c                 j   ^^ [         R                  " XX45      n[        U5      mU4S jmUU4S jnXV4$ )Nc                   > Uu  pVpvU R                  US9nU R                  U5         [        R                  " S5      " S5      n	U	n
US:X  a\  [	        ST5       HJ  nU R                  X[S-
  5      nU R                  X{5      nU R                  X5      nU R                  X5      n
ML     GO"SUs=:  a	  TS-
  :  a  O  O[	        SU5       HG  nU R                  X[5      nU R                  X{5      nU R                  X5      nU R                  X5      n
MI     [	        US-   T5       HJ  nU R                  X[S-
  5      nU R                  X{5      nU R                  X5      nU R                  X5      n
ML     OZ[	        STS-
  5       HG  nU R                  X[5      nU R                  X{5      nU R                  X5      nU R                  X5      n
MI     U R                  U5        S S S 5        X4$ ! , (       d  f       UW
4$ = f)Nname@   r   r_   )	append_basic_block
goto_blockr   IntTyperangeextract_valuemulre   branch)r4   	block_pos
block_namebb_endr~   rS   r6   rg   bbrM   flat_idxr  strideidx_imlen_idxs                  r)   	gen_block\DUFunc._install_ufunc_reduce.<locals>.ol_reduce.<locals>.compute_flat_idx.<locals>.gen_block-  s   )-&G 333DB ++B/!zz"~a0#'$>%*1g%6)0)>)>wA)N(/(=(=c(E$+KK$>+2;;x+C	 &7
 8Wq[8%*1i%8)0)>)>w)J(/(=(=c(E$+KK$>+2;;x+C	 &9 &+9q='%B)0)>)>wA)N(/(=(=c(E$+KK$>+2;;x+C	 &C &+1gk%:)0)>)>w)J(/(=(=c(E$+KK$>+2;;x+C	 &;  v.; 0> <'? 0/> x<'s   F4G&&
G7c                 F  > Uu  pEpgUR                   nUR                  SS9n	/ n
[        T5       H$  nT" XSU 3X5      u  pU
R                  X45        M&     UR	                  U5         UR                  XzS   S   5      n[        T5       H  nUR                  XU   S   5        M     S S S 5        UR                  U	5        UR                  U
S   S   R                  5      nU
 H  u  nnUR                  UU5        M     UR                  X5      $ ! , (       d  f       No= f)Naxis_endr  axis_r   r_   )basic_blockr  r  r  r  switchadd_caseposition_at_endphir   add_incomingsdiv)r3   r4   rh   r~   rS   itemsizerg   axisr  
switch_endlr  blockr  r  r  valuer  r  s                    r)   codegenZDUFunc._install_ufunc_reduce.<locals>.ol_reduce.<locals>.compute_flat_idx.<locals>.codegenR  s   370Gs ,,B!(!;!;!;!LJA"7^*3G%s4>+F%!23 ,
 !++B/!(eAh!?!&wA"OOAtAw7 "0 0
 ++J7!++ad1gll3C()u((6 )*"<<66 0/s   (AD
D )r   rE   rI   )	rS  rS   r  rg   r  rh   r  r  r  s	          @@r)   compute_flat_idxIDUFunc._install_ufunc_reduce.<locals>.ol_reduce.<locals>.compute_flat_idx(  s1    jjC>c(#(J7, |#r,   c                 ~    U n[        [        U 5      5       H"  nX   S:  a  X   U-   OX   n[        X#U5      nM$     U$ rv   )r  rI   r   )r  ndimaxr  r[   s        r)   
fixup_axisCDUFunc._install_ufunc_reduce.<locals>.ol_reduce.<locals>.fixup_axisj  sD    s4y)A,0GaK$'D.TWC&rc2B * 	r,   c                 h    SU S   p![        [        U 5      5       H  nX   U:  d  M  X0U   p!M     X4$ rv   )r  rI   )r  rg   r  r  s       r)   find_minADUFunc._install_ufunc_reduce.<locals>.ol_reduce.<locals>.find_minr  s:    CFQs3xAvz!"FQ ) vr,   c                    > T(       a#  Uc   [        U5      S:X  a  ST S3n[        U5      eSnT(       a  T
(       a  SnUS   nOT(       a  TnOUnUR                  S   n[        Xh5       H  n	U " XqU	   5      nM     U$ )Nr   'zero-size array to reduction operation  which has no identityr_   )rI   
ValueErrorrR   r  )r!   rW   r  ra   initialr!  startrr  r  id_noner   identity_none	init_none
ufunc_names             r)   impl_1d@DUFunc._install_ufunc_reduce.<locals>.ol_reduce.<locals>.impl_1dz  s     W_UqD(\)?AC$S/)EaA AA[[^u)Aaq*A *r,   c                   > Uc  [        S5      eUS:  a  X!R                  -  nUS:  d  X!R                  :  a  [        S5      eT(       a'  Uc$  UR                  U   S:X  a  ST S3n[        U5      eT" UR                  U5      nUcG  TcD  [        R                  " UTS9n[        R
                  " U5       H  u  pT" XS5      n
X   Xx'   M     O2Uc  Tb  [        R                  " UTTS9nO[        R                  " XdTS9nUR                  5       nUc[  TcX  [        R
                  " U5       H<  u  pX   S:X  a  M  T" UR                  UR                  X5      nX   UpU " X5      X'   M>     U$ [        R
                  " U5       HB  u  pUc  Tc
  X   S:X  a  M  T" UR                  UR                  X5      nX   UpU " X5      X'   MD     U$ )Nz'axis' must be specifiedr   Invalid axisr  r  )ra   )
fill_valuera   )
r  r  rR   r   emptyndenumeratefullravelrS   r  )r!   rW   r  ra   r  r!  rR   r  rg   r6   
result_idxviewr[   flat_poslhsrhsr  r   r  nb_dtyper  r  r  s                   r)   impl_nd_axis_intIDUFunc._install_ufunc_reduce.<locals>.ol_reduce.<locals>.impl_nd_axis_int  s   
 <$%?@@!8JJ&D!8tzz1$^44 W_T9Ja9OD(\)?AC$S/) $EKK6?x'7h7A"$.."3%71%E
!&!2 #4 _)= ((KAJA wwy?x'7$&NN5$99>$'7		1::8;(CH'+~s-23_DN %:"  %'NN5$9"?x/? #	Q$#3AIIqzz47$?#'>3S).s %: r,   c                 B  > T" X!R                   5      n[        S[        U5      5       H\  nXV   S:  d  XV   UR                   :  a  [        S5      e[        US-   [        U5      5       H  nXV   XW   :X  d  M  [        S5      e   M^     T" U5      u  pU R	                  UU	UUS9n
[        U5      S:X  a  U
$ [        U5      S:X  a  U R	                  XUS-   S-     S-
  S9$ Tn[        [        U5      5       H  nXh:w  d  M
  [        XXV   5      nM     U R	                  XS9$ )Nr   r  r_   zduplicate value in 'axis')r  ra   r  rF  )r  )r  r  rI   r  r  r   )r!   rW   r  ra   r  r  r  r  min_idxmin_elemr  r  axis_tupr  r  s               r)   impl_nd_axis_tupleKDUFunc._install_ufunc_reduce.<locals>.ol_reduce.<locals>.impl_nd_axis_tuple  s$   
 #44q#e*-Ax!|ux5::'=(88"1q5#e*5 8ux/",-H"II 6	 . %-UO!LL&.',)0 ! 2 t9>HY!^ <<w{a6G0H10L<MM!B"3r7^<!.reh!?B , !<<<33r,   c                     U$ r    r   )r!   rW   r  ra   r  s        r)   impl_axis_empty_tupleNDUFunc._install_ufunc_reduce.<locals>.ol_reduce.<locals>.impl_axis_empty_tuple  s	    
 r,   c                 *   > U R                  UTX45      $ r    )r  )r!   rW   r  ra   r  r  s        r)   impl_axis_noneGDUFunc._install_ufunc_reduce.<locals>.ol_reduce.<locals>.impl_axis_none  s    
 ||E8UDDr,   r   NN)r  r  r	   rl  r:   r   r  rm  UniTuplera   Integerro  rI   r   rr  r!   r   r   r   r  r   r   r   r  OmittedIntegerLiteral)r!   rW   r  ra   r  r!  axis_int_tupleaxis_empty_tuple	axis_noner  r  r  r  r  r  r  r  r  r  r   r  r  r  r  r  r  r  r  r(   s                 @@@@@@@@@@@@@@r)   	ol_reduce/DUFunc._install_ufunc_reduce.<locals>.ol_reduce  s    MM@#)#I#IK eU[[11E++C00'enn= 64::u}}5 )$<OTa++D1I JJ//47M,,J ##3t9q=.zl ;K K++C00uzz*H%**q.1K&-&9&9%&@&@u{{eH}}H))(3G++G4I
 
   ?$ ?$B     , '('+)-	< <@ )*)-+/	4 4@ ,-,0.2	 %&%)'+	E zzQ'7!,, uzz!23%%  3t9q=1))j05050D0D0F G G
 ('Gr,   r  r  )r(   r7  r8  r  s   `   r)   r>  DUFunc._install_ufunc_reduce  s,    ^^H%	X	&`	( 
'`	(r,   c                 |  > [         R                  " U5      (       a  U4OX4nS U 5       n[        S U 5       5      nU R                  U5      S:X  a2  U R                  (       a  SU < S3n[        U5      eU R                  " U6   [         R                  " U5      (       a  [        TU ]!  X5      $ [        TU ]   " XU46 $ )Nc              3   8   #    U  H  n[        U5      v   M     g 7fr    )r   r  r   s     r)   r  DUFunc.at.<locals>.<genexpr>  s     .#&++r  c              3   |   #    U  H2  n[        U[        R                  5      (       a  UR                  OUv   M4     g 7fr    )r:   r   r  ra   r  s     r)   r  r     s0      /'- *4C)E)ECII3N'-s   :<r   r  z.at(...))	r   rr  r   r   r   r  r  r   r8  )	r(   r"   r$   r&   r~   r  r'  r!  r   s	           r)   r8  	DUFunc.at  s    **1--tA6.. /'-/ / ##K0L@|||?CE"3''""D) q!!7:a))7:A//r,   c           	      $   Uc   U R                   R                  R                  n[        SU R                  R
                  -   [        4[        X R                  S95      nUR                  X5        U R                  U5        U R                  U5        g)a  Constructs and installs a typing class for a DUFunc object in the
input typing context.  If no typing context is given, then
_install_type() installs into the typing context of the
dispatcher object (should be same default context used by
jit() and njit()).
NDUFuncTyping_)rw   generic)r   targetdescrtyping_contextr   r!   r   r   r   _type_meinsert_user_functionr;  r@  )r(   rS  _ty_clss      r)   r   DUFunc._install_type"  sz     ((44CCI)<)<<(*mm<> 	&&t5&&w/##G,r,   c                 :   U R                   (       aH  [        R                  " X5      nUc  g[        UR                  UR
                  -   5      S[        U5       nU R                  R                  R                  5        H  u  p4UR                  U:X  d  M  X44s  $    g)z
Given a tuple of element-wise argument types, find a matching
signature in the dispatcher.

Return a 2-tuple containing the matching signature, and
compilation result.  Will return two None's if no matching
signature was found.
Nr   )r   r   ufunc_find_matching_loopr   inputsoutputsrI   r   r   r  r~   )r(   r'  looprh   rL   s        r)   r   DUFunc.find_ewise_function2  s     << 99$LD|!dll :;<MS=MNK))3399;ICxx;&y  < r,   c                    U(       a   eU R                   n[        R                  R                  X1U5      nUu  pVpx[	        U5      n	U	S:  a  [        US[	        U5      *  5      n
O[        U5      n
U R                  U
5      u  pUcX  U R                  (       a  [        R                  " SU < SU< 35      eU R                  U
5        U R                  U
5      u  pUc   eU	S:  a  [        U5      nO]UR                  S:X  a7  US:  a#  [        R                  " UR                  Xx5      /nO$UR                  /nO[        R                   " S5      eUR#                  U5        [%        U6 $ )z
Implement AbstractTemplate.generic() for the typing class
built by DUFunc._install_type().

Return the call-site signature after either validating the
element-wise signature or compiling for it.
r   Nzcannot call z with types r_   ztyping gufuncs (nout > 1))r!   r   Numpy_rules_ufunc_handle_inputsrI   r   r   r   r	   rm  r   r   r   r   r  ro   NumbaNotImplementedErrorextendr   )r(   r  kwtysr!   _handle_inputs_result
base_typesexplicit_outputsndimsrV   explicit_output_countr'  rh   rL   outtyss                 r)   r  DUFunc._type_meF  sL    y

 ' 9 9 H H5!"6K3
e #$4 5 1$
+BS1A-B,B CDK
+K,,[9	; ||++/3V-= > >$$[100=IC?"? 1$*+FZZ1_qy++coouEF//*112MNNf&!!r,   )r   r   r   r   r   r    )r,  N)'r   r   r   r   r   set_DUFunc__base_kwargsr*   r   r   classmethodr   r   propertyr   r   r   r   r   r   r   r   r   re   r   r  r   r   r(  r;  r@  r?  r>  r8  r   r   r  r   r   r   s   @r)   r   r      sC    ABM)-U" 0


   . .         ! !     # # $ $;27> " "H'2*)c$Jd(L	0*- (&" &"r,   r   )3r   rc   r  numpyr   numbar   r   
numba.corer   r   r   r   r	   numba.core.extendingr
   r   r   r   r   numba.core.typingr   numba.core.typing.templatesr   r   numba.cpython.unsafe.tupler   numba.np.ufuncr   numba.np.ufunc.ufunc_baser   r   numba.parforsr   r   r   r   typingr   llvmliter   numba.core.compiler_lockr   r   r   r   ReduceMixinr   r   	MAP_TYPESr  r   r,   r)   <module>r2     s         B B- - & C 4 $ A ( ' "   9j= j=Z(5$ 5H"Y""I$5$5y H"V      'r,   