o
    2h                     @   s	  d dl Z d dlZd dlm  mZ d dlmZ d dl	m
Z
 d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d d	lmZ d d
lmZ duddZejddddd Zdd Zdd Zdd Zdd Zdd ZdvddZd d! Z ejdddd"d# Z!d$d% Z"d&d' Z#dwd(d)Z$dxd*d+Z%dyd,d-Z&dyd.d/Z'ej(dd0d1d2 Z)d3d4 Z*dwd5d6Z+d7d8 Z,dwd9d:Z-dwd;d<Z.dwd=d>Z/dyd?d@Z0dzdAdBZ1ej2dCdD Z3ej2dEdF Z4ej(dd0dGdH Z5ej(dd0dIdJ Z6ej(dd0dKdL Z7dMdN Z8ej(dd0dOdP Z9dwdQdRZ:dwdSdTZ;d{dVdWZ<dydXdYZ=d|dZd[Z>d\d] Z?d^d_ Z@d`da ZAdbdc ZBddde ZCdfdg ZDdhdi ZEdjdk ZFdldm ZGdndo ZHdpdq ZIdrds ZJej(dd0dtdu ZKdvdw ZLd}dxdyZMej(dd0dzd{ ZNej(dd0d|d} ZOej(dd0d~d ZPej2dd ZQej2dd ZRdyddZSd~ddZTd|ddZUd|ddZVdd ZWd}ddZXd}ddZYdddZZdddZ[ej(dd0dd Z\dd Z]dyddZ^dd Z_ej2dd Z`ej2dd Zadd Zbej(dd0dd ZcdyddZdej(dd0dd ZedyddZfdyddZgdd Zhdd Zidd ZjdyddZkej(dd0dd ZldddZmej2dd Znej(dd0dd Zoej(dd0dd ZpddÄ Zqddń Zr	 dddȄZsej2ddʄ Ztej2dd̄ Zuej(dd0dd΄ Zvej2ddЄ Zwdd҄ ZxddԄ Zyddք Zzdd؄ Z{dddۄZ|ejddddd݄ Z}dwdd߄Z~ddddZdxddZejddddd Zdd Zdd ZdddZdd Zdd Zdd ZdyddZdyddZdd ZdddZdxddZdddZdd Zdd Zej(dd0dd Zej2d	d
 ZdyddZdd ZdyddZdddZej(dd0dd Zej(dd0dd Zej(dd0dd Zej(dd0dd Zdd Zd{ddZd}d d!Zd}d"d#Zdwd$d%Zd&d' Zdyd(d)Zdyd*d+Zej(dd0d,d- Zej(dd0d.d/ Zdd1d2Zej(dd0d}d3d4Zd5d6 Zdd7d8Zdd9d:Zd}d;d<Zd}d=d>Zd?d@ ZdAdB ZdCdD ZdEdF ZdddGdHdIZd|dJdKZejdLdM ZdNdO ZdPdQ ZdRdS Zej(dd0dTdU Zej(dd0dVdW Zej(dd0dXdY ZdydZd[Zdyd\d]Zdwd^d_Zdwd`daZddbdcZddde Zdfdg Zdhdi ZddkdlZd}dmdnZdodp ZÐd{dqdrZĐddsdtZdS (      N)config)dtypes)canonicalize_axis)to_tuple_or_list)standardize_dtype)nnsparse)castconvert_to_tensor   r   r   c                 C   sV   | j dk rtd| j  t|dks|d |d kr#td| dtj| ||dS )z5Rotate an array by 90 degrees in the specified plane.   zBInput array must have at least 2 dimensions. Received: array.ndim=r   r   zInvalid axes: z3. Axes must be a tuple of two different dimensions.)kaxes)ndim
ValueErrorlenjnprot90)arrayr   r    r   V/var/www/html/chatgem/venv/lib/python3.10/site-packages/keras/src/backend/jax/numpy.pyr      s   

r   T)linearuse_sparsifyc                 C      t | } t |}t| |S N)r   r   addx1x2r   r   r   r   !      r   c                 C      t | } t| S r   )r   r   bartlettxr   r   r   r$   (      
r$   c                 C   r#   r   )r   r   hammingr%   r   r   r   r(   -   r'   r(   c                 C   r#   r   )r   r   hanningr%   r   r   r   r)   2   r'   r)   c                 C   r   r   )r   r   	heavisider   r   r   r   r*   7      r*   c                 C      t | } t| |S r   )r   r   kaiser)r&   betar   r   r   r-   =      r-   Fc                    sJ  |st | tjrot | tjr7|d ur4t |tjstd| j|jur1tt| j|js1td|j}| j} t	| j
dkr@dnd}tt| d  }tj| |dd}|d uret||d }|| }tj||fd}|S t	| j
dkr|d u r fd	d
}	tt|	| }
n fdd
}	tt|	t| |}
t|
S tj| | dS )Nz$`x` and `weights` must both be BCOOsz2`x` and `weights` BCOOs must have the same indicesr   r   Tr   r   r   c                    s   t j|  dS )N	minlengthr   bincount)arrr1   r   r   bincount_fna      zbincount.<locals>.bincount_fnc                    s   t j| d | d  dS )Nr   r   weightsr2   r3   )arr_wr1   r   r   r6   g   s   r8   )
isinstance
jax_sparseBCOOr   indicesr   allequaldatar   shapemaximummaxr   one_hotexpand_dimsbcoo_reduce_sumlistmapzipstackr4   )r&   r9   r2   r	   reduction_axis	maxlengthone_hot_encodingexpanded_weightsoutputsr6   	bincountsr   r1   r   r4   B   s>   
r4   c                 O   sf   dd |D }t tdd |D }t|dkr!|d dkr!d}nd }||d	< tj| g|R i |S )
Nc                 S   s   g | ]}t |qS r   r   .0r&   r   r   r   
<listcomp>s   s    zeinsum.<locals>.<listcomp>c                 s   s    | ]}t |jV  qd S r   )r   dtyperR   r   r   r   	<genexpr>v       zeinsum.<locals>.<genexpr>r   r   int8int32preferred_element_type)rH   setr   r   einsum)
subscriptsoperandskwargsr   rZ   r   r   r   r\   r   s   r\   c                 C   r   r   )r   r   subtractr   r   r   r   r`      r"   r`   c                 C   s   t | } t |}t| j}t|j}|dkr|dkrd}nd }t| tjs+t|tjrettds7tt	jt_
t| tjrJtj| t| jd dd} t|tjr]tj|t|jd dd}tj
| ||dS t	j| ||dS )NrX   rY   sparse_matmulr   warn)n_batchon_inefficient)rZ   )r   r   rU   r;   r<   	JAXSparsehasattrmatmulsparsifyr   ra   r=   bcoo_update_layoutr   rB   )r    r!   x1_dtypex2_dtyperZ   r   r   r   rg      s0   


rg   c                 C   s   t | } t |}t| tjrUt|tjr@| j|ju r:| js't| } t|}tjt| j	|j	| jf| j
dddS t| |S t| |}tj|| jf| j
| j| jdS t|tjrpt|| }tj||jf|j
|j|jdS t| |S )NT)rB   indices_sortedunique_indices)r   r;   r<   r=   r>   rm   bcoo_sum_duplicatesr   multiplyrA   rB   bcoo_multiply_sparsebcoo_multiply_denserl   )r    r!   out_datar   r   r   ro      s>   

ro   c                    s   t   t j}t jd}d|v s|dkr|}n|}t tjrl|d u r0tt	t
 j}tj| jdd\}}}t fdd|D }	tj |d}
tj|
j||	 |
jf|
jd	}
|rjtj|
||d
}
|
S tj |||d}
t|
|S )Nfloat32intboolFinsert_dimsc                 3   s    | ]} j | V  qd S r   rB   )rS   ir%   r   r   rV      rW   zmean.<locals>.<genexpr>r0   rx   rB   broadcast_dimensionsaxiskeepdimsrU   )r   r   rU   r   result_typer;   r<   r=   tupleranger   rB   r	   $axis_shape_dims_for_broadcast_in_dimmathprodrG   rA   astyper>   bcoo_broadcast_in_dimr   meanr
   )r&   r}   r~   	ori_dtypecompute_dtyperesult_dtypecanonical_axiskeep_dims_shaper{   divisoroutputr   r%   r   r      s>   

r   c                 C      t | } tj| |||dS N)r}   r~   initial)r   r   rD   r&   r}   r~   r   r   r   r   rD         rD   c                 C      |pt  }tj| |dS N)rU   )r   floatxr   onesrB   rU   r   r   r   r         r   c                 C   r   r   )r   r   r   zerosr   r   r   r   r      r   r   )r   c                 C   r#   r   )r   r   absoluter%   r   r   r   r        
r   c                 C      t | S r   )r   r%   r   r   r   abs     r   c                 C      t j| ||dS Nr}   r~   )r   r?   r&   r}   r~   r   r   r   r?        r?   c                 C   B   t | } t| jdkrt }nt| jt}t| |} t	
| S Nint64)r   r   rU   r   r   r   r   floatr
   r   angler&   rU   r   r   r   r     s   


r   c                 C   r   r   )r   anyr   r   r   r   r     r   r   c                 C   r   r   )r   amaxr   r   r   r   r   !  r   r   c                 C   r   r   )r   aminr   r   r   r   r   %  r   r   c                 C       t | } t |}tj| ||dS Nr}   )r   r   append)r    r!   r}   r   r   r   r   )     r   c                 C   sj   |d u r(t | dt| t |dt|g}|d ur#|t |dt| tj| }t|}tj| |||dS )NrU   )steprU   )getattrtyper   r   r   r   r   arange)startstopr   rU   dtypes_to_resolver   r   r   r   /  s   
r   c                 C   r   r   )r   r   rU   r   r   r   r   r   r
   r   arccosr   r   r   r   r   <     


r   c                 C   r   r   )r   r   rU   r   r   r   r   r   r
   r   arccoshr   r   r   r   r   G  r   r   c                 C   r   r   )r   r   rU   r   r   r   r   r   r
   r   arcsinr   r   r   r   r   R  r   r   c                 C   r   r   )r   r   rU   r   r   r   r   r   r
   r   arcsinhr   r   r   r   r   ]  r   r   c                 C   r   r   )r   r   rU   r   r   r   r   r   r
   r   arctanr   r   r   r   r   h  r   r   c                 C   B   t | } t |}t| j|jt}t| |} t||}t| |S r   )r   r   r   rU   r   r
   r   arctan2r    r!   rU   r   r   r   r   s     

r   c                 C   r   r   )r   r   rU   r   r   r   r   r   r
   r   arctanhr   r   r   r   r   |  r   r   c                 C      ddl m} t| } t| j}d|vs| r| jdkr#tj| ||dS t	|d}t
| |} | dkt| @ }t|t| jj | } tj| ||dS Nr   )uses_cpur   r   rs           )keras.src.testing.test_caser   r   r   rU   r   r   argmaxr   r   r
   signbitwherefinfotinyr&   r}   r~   r   rU   is_negative_zeror   r   r   r        

r   c                 C   r   r   )r   r   r   r   rU   r   r   argminr   r   r
   r   r   r   r   r   r   r   r   r     r   r   c                 C   s.   t | } | jdkrtj| d dS tj| |dS )Nr   r   )r   r   r   argsortr&   r}   r   r   r   r     s   
r   c                 C      t j| |dS r   )r   r   r   r   r   r   r     r7   r   c                 C   sd   t | } | jtg}|d urt |}||j tj| }t| |} |d ur*t||}tj| ||dS )N)r9   r}   )	r   rU   r   r   r   r   r
   r   average)r&   r}   r9   r   rU   r   r   r   r     s   



r   c                 C   r   r   )r   r   bitwise_andr&   yr   r   r   r     r+   r   c                 C   r#   r   )r   r   invertr%   r   r   r   bitwise_invert  r'   r   c                 C   r   r   )r   r%   r   r   r   bitwise_not  r   r   c                 C   r   r   )r   r   
bitwise_orr   r   r   r   r     r+   r   c                 C   r   r   )r   r   bitwise_xorr   r   r   r   r     r+   r   c                 C   &   t | } t|tst |}t| |S r   )r   r;   rt   r   
left_shiftr   r   r   r   bitwise_left_shift     
r   c                 C   
   t | |S r   )r   r   r   r   r   r        
r   c                 C   r   r   )r   r;   rt   r   right_shiftr   r   r   r   bitwise_right_shift  r   r   c                 C   r   r   )r   r   r   r   r   r     r   r   c                 C   r#   r   )r   r   blackmanr%   r   r   r   r     r'   r   c                 C   r,   r   )r   r   broadcast_to)r&   rB   r   r   r   r     r/   r   c                 C   r#   r   )r   r   cbrtr%   r   r   r   r     r'   r   c                 C   r   r   )r   r   rU   r   r   r   r   r   r
   r   ceilr   r   r   r   r     r   r   c                 C   s.   t | } t| jdkrt| d} t| ||S )Nru   rY   )r   r   rU   r
   r   clip)r&   x_minx_maxr   r   r   r     s   
r   c                 C   sb   t dd | D }|r*|t| kr#t|t| d j}tj| |dS dd | D } tj| |dS )Nc                 s   s    | ]	}t |tjV  qd S r   )r;   r<   r=   rR   r   r   r   rV     s    zconcatenate.<locals>.<genexpr>r   )	dimensionc                 S   s$   g | ]}t |tjr| n|qS r   )r;   r<   re   todenserR   r   r   r   rT     s    zconcatenate.<locals>.<listcomp>r   )	builtinssumr   r   rB   r<   bcoo_concatenater   concatenate)xsr}   
bcoo_countr   r   r   r     s   r   c                 C   r#   r   r   r   	conjugater%   r   r   r   r   !  r   r   c                 C   r#   r   r   r%   r   r   r   conj'  r   r   c                 C   r#   r   )r   r   copyr%   r   r   r   r   -  r   r   c                 C   r   r   )r   r   rU   r   r   r   r   r   r
   r   cosr   r   r   r   r   3  r   r   c                 C   r   r   )r   r   rU   r   r   r   r   r   r
   r   coshr   r   r   r   r   >  r   r   c                 C   s   t tj| |ddS )Nr   rY   )r
   r   count_nonzeror   r   r   r   r   I  s   r   c                 C   s&   t | } t |}tj| |||||dS )N)axisaaxisbaxiscr}   )r   r   cross)r    r!   r   r   r   r}   r   r   r   r  M  s   r  c                 C      t | } tj| ||dS N)r}   rU   )r   r   cumprodr&   r}   rU   r   r   r   r  Z     r  c                 C   r  r  )r   r   cumsumr  r   r   r   r  _  r  r  c                 C   r#   r   )r   r   deg2radr%   r   r   r   r  d  r'   r  c                 C      t | } tj| |dS N)r   )r   r   diagr&   r   r   r   r   r  i     r  c                 C   r	  r
  )r   r   diagflatr  r   r   r   r  n  r  r  c                 C   r   )N)offsetaxis1axis2)r   r   diagonal)r&   r  r  r  r   r   r   r  s  s   r  c                 C   r  )N)nr}   )r   r   diff)ar  r}   r   r   r   r  }  r  r  c                 C   r   r   )r   r   digitize)r&   binsr   r   r   r    r"   r  c                 C   r   r   )r   r   dotr   r   r   r   r    r+   r  c                 C   r   r   )r   r   r   emptyr   r   r   r   r    r   r  c                 C   r   r   )r   r   r@   r   r   r   r   r@     r+   r@   c                 C   :   t | } t| j}d|v s|dkrt| t } t| S Nrt   ru   )r   r   rU   r
   r   r   r   expr&   r   r   r   r   r    
   

r  c                 C   r  r  )r   r   rU   r
   r   r   r   exp2r  r   r   r   r    r  r  c                 C   sH   t | } t| tjrtj|| jdd\}}}tj| ||dS t	| |S )NTrv   rz   )
r   r;   r<   r=   r	   r   rB   r   r   rF   )r&   r}   _result_shaper{   r   r   r   rF     s   rF   c                 C   r  r  )r   r   rU   r
   r   r   r   expm1r  r   r   r   r"    r  r"  c                 C   r   r   )r   flipr   r   r   r   r#    r7   r#  c                 C   r   r   )r   r   rU   r   r   r   r   r   r
   r   floorr   r   r   r   r$    r   r$  c                 C   s   |pt  }tj| ||dS r   )r   r   r   full)rB   
fill_valuerU   r   r   r   r%    s   r%  c                 C   r   r   )r   	full_like)r&   r&  rU   r   r   r   r'    r   r'  c                 C   r   r   )r   r   greaterr   r   r   r   r(    r+   r(  c                 C   r   r   )r   r   greater_equalr   r   r   r   r)    r+   r)  c                 C   
   t | S r   )r   hstackr   r   r   r   r+    r   r+  c                 C   r   r   )r   r   r   identity)r  rU   r   r   r   r-    r   r-  c                 C   r#   r   )r   r   imagr%   r   r   r   r.    r   r.  h㈵>:0yE>c                 C   s"   t | } t |}t| ||||S r   )r   r   isclose)r    r!   rtolatol	equal_nanr   r   r   r1    s   r1  c                 C   r#   r   )r   r   isfiniter%   r   r   r   r5    r   r5  c                 C   r#   r   )r   r   isinfr%   r   r   r   r6    r   r6  c                 C   r#   r   )r   r   isnanr%   r   r   r   r7  
  r   r7  c                 C   r   r   )r   r   lessr   r   r   r   r8    r+   r8  c                 C   r   r   )r   r   
less_equalr   r   r   r   r9    r+   r9  2   c              	   C      t j| ||||||dS )N)numendpointretsteprU   r}   )r   linspace)r   r   r<  r=  r>  rU   r}   r   r   r   r?    s   r?  c                 C   .   t | } t| jdkrt| t } t| S r   )r   r   rU   r
   r   r   r   logr%   r   r   r   rA  *     
rA  c                 C   r@  r   )r   r   rU   r
   r   r   r   log10r%   r   r   r   rC  2  rB  rC  c                 C   r@  r   )r   r   rU   r
   r   r   r   log1pr%   r   r   r   rD  :  rB  rD  c                 C   r@  r   )r   r   rU   r
   r   r   r   log2r%   r   r   r   rE  B  rB  rE  c                 C   r   r   )r   r   r   rU   r   r
   r   	logaddexpr   r   r   r   rF  J  r   rF  c                 C   r   r   )r   r   logical_andr   r   r   r   rG  S  r+   rG  c                 C   r#   r   )r   r   logical_notr%   r   r   r   rH  Y  r'   rH  c                 C   r   r   )r   r   
logical_orr   r   r   r   rI  ^  r+   rI  
   c              	   C   r;  )N)r<  r=  baserU   r}   )r   logspace)r   r   r<  r=  rK  rU   r}   r   r   r   rL  d  s   rL  c                 C   r   r   )r   r   rC   r   r   r   r   rC   p  r"   rC   c                 C   s   t |tr	t|}t| } t| jdkrt| t } t	j
| ||d}|du r>|d u r>|j| jk r>t	j|dd}|j| jk s1|S )Nr   r   Tr   r   )r;   rH   r   r   r   rU   r
   r   r   r   medianr   rF   )r&   r}   r~   resultr   r   r   rM  w  s   
rM  xy)indexingc                 G   s   t j|d| iS )NrP  )r   meshgrid)rP  r&   r   r   r   rQ    r   rQ  c                 C   r   r   )r   r   minr   r   r   r   rR    r   rR  c                 C   r   r   )r   r   minimumr   r   r   r   rS    r"   rS  c                 C   r   r   )r   r   modr   r   r   r   rT    r+   rT  c                 C   r   )N)sourcedestination)r   moveaxis)r&   rU  rV  r   r   r   rW    r   rW  r   c                 C   r   )N)nanposinfneginf)r   r   
nan_to_num)r&   rX  rY  rZ  r   r   r   r[    r   r[  c                 C   r*  r   )r   r   r%   r   r   r   r     r   r   c                 C   r*  r   )r   nonzeror%   r   r   r   r\    r   r\  c                 C   r   r   )r   r   	not_equalr   r   r   r   r]    r+   r]  c                 C   r   r   )r   	ones_liker   r   r   r   r^    r7   r^  c                 C   r   r   )r   
zeros_liker   r   r   r   r_    r7   r_  c                 C      t | |S r   )r   outerr   r   r   r   ra       ra  constantc                 C   sJ   t | } i }|d ur|dkrtd| ||d< tj| |fd|i|S )Nrc  zZArgument `constant_values` can only be provided when `mode == 'constant'`. Received: mode=constant_valuesmode)r   r   r   pad)r&   	pad_widthre  rd  r_   r   r   r   rf    s   rf  c                 C   r   )Nr|   )r   r   r   )r&   r}   r~   rU   r   r   r   r     r   r   r   c                 C   s   t | } t |}t| jdkrt| t } tj| ||||d}|du rG|d u rG| jt	|j
dkr3dnd }|j|k rGtj|dd}|j|k s;|S )Nr   )r}   methodr~   Tr   r   r   r   )r   r   rU   r
   r   r   r   quantiler   r   rB   rF   )r&   qr}   rh  r~   rN  result_ndimr   r   r   ri    s   

ri  c                 C   r#   r   )r   r   ravelr%   r   r   r   rl    r'   rl  c                 C   r,   r   )r   r   unravel_index)r>   rB   r   r   r   rm    r/   rm  c                 C   r#   r   )r   r   realr%   r   r   r   rn    r   rn  c                 C   r#   r   )r   r   
reciprocalr%   r   r   r   ro    r   ro  c                 C   r  r   )r   r   repeat)r&   repeatsr}   r   r   r   rp    r  rp  c                 C   sN   t | tjr!ddlm} || j|d}d |vr|}tj| |dS t	| |S )Nr   )operation_utils	new_shape)	new_sizes)
r;   r<   r=   keras.src.opsrr  compute_reshape_output_shaperB   bcoo_reshaper   reshape)r&   newshaperr  output_shaper   r   r   rx    s   rx  c                 C   r   r   )r   roll)r&   shiftr}   r   r   r   r{    r   r{  leftc                 C   s,   t | dkrtd| j tj| ||dS )Nr   z`searchsorted` only supports 1-D sorted sequences. You can use `keras.ops.vectorized_map` to extend it to N-D sequences. Received: sorted_sequence.shape=)side)r   r   rB   r   searchsorted)sorted_sequencevaluesr~  r   r   r   r    s   r  c                 C   r#   r   )r   r   signr%   r   r   r   r    r   r  c                 C   r#   r   )r   r   r   r%   r   r   r   r   "  r   r   c                 C   r   r   )r   r   rU   r   r   r   r   r   r
   r   sinr   r   r   r   r  (  r   r  c                 C   r   r   )r   r   rU   r   r   r   r   r   r
   r   sinhr   r   r   r   r  3  r   r  c                 C   r*  r   )r   sizer%   r   r   r   r  >  r   r  c                 C   r	  r   )r   r   sortr   r   r   r   r  B  r  r  c                 C   r   r   )r   split)r&   indices_or_sectionsr}   r   r   r   r  G  r   r  c                 C   r   r   )r   rK   r   r   r   r   rK   K  r7   rK   c                 C   s4   t | } t| jdkrt| t } tj| ||dS )Nr   r   )r   r   rU   r
   r   r   r   stdr   r   r   r   r  O  s   r  c                 C   r  )N)r  r  )r   r   swapaxes)r&   r  r  r   r   r   r  V  r  r  c                 C   s$   t | } t |dd}tj| ||dS )NFr   r   )r   r   taker&   r>   r}   r   r   r   r  [  s   r  c                 C   r   r   )r   take_along_axisr  r   r   r   r  a  r   r  c                 C   r   r   )r   r   rU   r   r   r   r   r   r
   r   tanr   r   r   r   r  e  r   r  c                 C   r   r   )r   r   rU   r   r   r   r   r   r
   r   tanhr   r   r   r   r  p  r   r  r   c                 C   r   )Nr0   )r   r   	tensordot)r    r!   r   r   r   r   r  {  r   r  c                 C   s   t | } t| j}d|v r9|dk r9ttd|t }t| t } t	| |} t
| } t| |} t| |S tj
| |dS )Nrt   r   rJ  )decimals)r   r   rU   r
   r   powr   r   r   ro   rounddivide)r&   r  x_dtypefactorr   r   r   r    s   


r  c                 C   r`  r   )r   tile)r&   rq  r   r   r   r    rb  r  c                 C   s2   t | } d }t| jdv rd}tj| ||||dS )N)ru   uint8uint16rY   )r  r  r  rU   )r   r   rU   r   trace)r&   r  r  r  rU   r   r   r   r    s
   r  c                 C      |pt  }tj| |||dS N)Mr   rU   )r   r   r   triNr  r   rU   r   r   r   r       r  c                 C   r	  r
  )r   r   trilr  r   r   r   r    r  r  c                 C   r	  r
  )r   r   triur  r   r   r   r    r  r  c                 C   s0   t | } t| j}d|v sd|kr| S t| S r  )r   r   rU   r   truncr   r   r   r   r    s
   

r  c                 C   r   r   )r   r   vdotr   r   r   r   r    r+   r  c                 C   r   r   )r   r   innerr   r   r   r   r    r+   r  c                 C   r*  r   )r   vstackr,  r   r   r   r    r   r  excluded	signaturec                C   s   |d u rt  }tj| ||dS )Nr  )r[   r   	vectorize)pyfuncr  r  r   r   r   r    s   r  c                 C      t | ||S r   )r   r   )	conditionr    r!   r   r   r   r     r7   r   c                 C   r   r   )r   r   r  r   r   r   r   r    r"   r  c                 C   s<   t | } t |}t|dkd|}t|dkdt| |S )Nr   r   )r   r   r   r  )r    r!   safe_x2r   r   r   divide_no_nan  s   r  c                 C   r   r   )r  r   r   r   r   true_divide  r   r  c                 C   r   r   )r   r   powerr   r   r   r   r    r+   r  c                 C   r#   r   )r   r   negativer%   r   r   r   r    r   r  c                 C   r#   r   )r   r   squarer%   r   r   r   r    r   r  c                 C   r@  r   )r   r   rU   r
   r   r   r   sqrtr%   r   r   r   r    rB  r  c                 C   sP   t | tjr!|d u rtdd t| jD }t|}tj| |dS tj	| |dS )Nc                 s   s     | ]\}}|d kr|V  qdS )r   Nr   )rS   ry   dr   r   r   rV     s    zsqueeze.<locals>.<genexpr>)
dimensionsr   )
r;   r<   r=   r   	enumeraterB   r   bcoo_squeezer   squeezer   r   r   r   r    s   r  c                 C   s|   t | } t| tjr7t| j}|d u rtt|d d d }ng }|D ]}t||}|	| q#tj
| |dS tj| |dS )Nr   )permutationr0   )r   r;   r<   r=   r   rB   r   r   r   r   bcoo_transposer   	transpose)r&   r   num_dimsr  r  r   r   r   r    s   

r  c                 C   s<   t | } t| jd}t| jt}ttj| |||d|S )Nrs   r|   )r   r   r   rU   r   r
   r   var)r&   r}   r~   r   r   r   r   r   r    s   r  c                 C   s|   t | } t| tjr6|d u rttt| j}tj	|| jdd\}}}tj
| |d}|r4tj|||d}|S tj| ||dS )NFrv   r0   rz   r   )r   r;   r<   r=   r   r   r   rB   r	   r   rG   r   r   r   )r&   r}   r~   r   r   r{   r   r   r   r   r   #  s(   r   c                 C   r  r  )r   r   r   eyer  r   r   r   r  <  r  r  c                 C   r   r   )r   r   floor_divider   r   r   r   r  A  r+   r  c                 C   r   r   )r   r   logical_xorr   r   r   r   r  G  r+   r  c                 C   r#   r   )r   r   corrcoefr%   r   r   r   r  M  r'   r  validc                 C   s   t | } t |}t| ||S r   )r   r   	correlate)r    r!   re  r   r   r   r  R  s   r  c                 C   r   )N)default)r   select)condlist
choicelistr  r   r   r   r  X  r   r  c                 C   s   t | } ttj| S r   )r   r   r   linalgslogdetr%   r   r   r   r  \  r  r  c                 C   r  r   )r   argpartition)r&   kthr}   r   r   r   r  a  r7   r  c                 C   r   )N)r  r   )r   	histogram)r&   r  r   r   r   r   r  e  r   r  )r   r   )Nr   F)NF)NFNr   )Nr   N)r   )NN)r   )r   r   r   N)r   r   r   )r   r   )r/  r0  F)r:  TFNr   )r:  TrJ  Nr   )r   NN)rc  N)Nr   F)r}  )r   )Nr   N)r  )rJ  N)r   r   jax.experimental.sparseexperimentalr	   r<   	jax.numpynumpyr   keras.src.backendr   keras.src.backend.commonr   &keras.src.backend.common.backend_utilsr   r   "keras.src.backend.common.variablesr   keras.src.backend.jaxr   keras.src.backend.jax.corer
   r   r   elementwise_binary_unionr   r$   r(   r)   r*   r-   r4   r\   r`   rg   ro   r   rD   r   r   elementwise_unaryr   r   r?   r   r   r   r   r   r   densifying_unaryr   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r  r  r  r  r  r  r  r  r  r  r  r@   r  r  rF   r"  r#  r$  r%  r'  r(  r)  r+  r-  r.  r1  r5  r6  r7  r8  r9  r?  rA  rC  rD  rE  rF  rG  rH  rI  rL  rC   rM  rQ  rR  rS  rT  rW  r[  r   r\  r]  r^  r_  ra  rf  r   ri  rl  rm  rn  ro  rp  rx  r{  r  r  r   r  r  r  r  r  rK   r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r   elementwise_divisionr  r  r  r  r  r  r  r  r  r  r   r  r  r  r  r  r  r  r  r  r   r   r   r   <module>   s   
0
)(
















	




































	
















	