o 2h7@s<ddlZddlZddlmZddlmZddlmZddlmZddlm Z ddlm Z ddl m Z d d Z d;d d Zd;ddZdd&d'Zd>d(d)Z d?d+d,Z  * d@d-d.Zd/d0Zd1d2Zd3d4Z d5d6Z!dAd7d8Z"d9d:Z#dS)BN)config)standardize_dtype)dtypes)cast)convert_to_tensor) get_device)padcCsttj|jddtd}||j|jtj }|p't t |}t |dk||}t ||k||}|dft |jdd}|dkr[tj|dtitd }n tj|dti}|d|||}|dddf}||jS) N)deviceramaxr Inf.)torchprodtensorshaperlongrepeat_interleaveviewtypeint64lenuniquewheretupleonesfloatzerosscatter_reducedtype)data segment_idsreduction_method num_segments num_repeatsrresultr&W/var/www/html/chatgem/venv/lib/python3.10/site-packages/keras/src/backend/torch/math.py_segment_reduction_fns0   r(FcCt|}t|}t||d|S)Nsumrr(r r!r#sortedr&r&r' segment_sum7r.cCr))Nr r+r,r&r&r' segment_max=r/r0TcCst|}tj|||dS)N)r-)rrtopk)xkr-r&r&r'top_kCsr4cCsZt|tj}|dddf}t|}t||j}tj||dd}||k}tj|ddS)Nr dim)axis)rrrrr4valuestake_along_dimany)targets predictionsr3 topk_valuestargets_valuesmaskr&r&r'in_top_kHs r@cCs4t|}|durtt|n|}tj|||dS)N)r6keepdim)rrranger6r logsumexp)r2r7keepdimsr&r&r'rCRsrCreducedcCs6t|}|dvrtd|t|}tjj||dS)N>rEcompletez]`mode` argument value not supported. Expected one of {'reduced', 'complete'}. Received: mode=)mode)r ValueErrorrlinalgqr)r2rGr&r&r'rJXsrJcCst|}tj|d||dS)Nr ) dimensionsizestep)rr unfold_copy)r2sequence_lengthsequence_strider&r&r'extract_sequencesdsrQc Cs~t|}|j^}}}||krtd|d||||kr,td|d|d|t|}t||||f}||d|}d|d|}||}tjj |d||ddddf}t|||||f}t |d}tjj |ddd|ddddf}|jdd} t||d f}|ddd|| |f}t|||| |f}tj |dd ddd|f}t|t |d S) NzQ`sequence_stride` must equal or less than x.shape[-1]. Received: sequence_stride=z, x.shape[-1]=zb`sequence_stride` must equal or greater than x.shape[-1] / x.shape[-2]. Received: sequence_stride=z, x.shape[-2]=r r)rr rRr r5)r ) rrrHmathrrreshapenn functionalrpermuter*r) r2rP batch_shape num_sequencesrOflat_batchsize output_sizenstep_per_segmentpadded_segment_lenshrinkedr&r&r'_overlap_sequencesksR     r`cCst|ttfr t|dkrtd||\}}t|}t|}|j|jkr2td|jd|jt|rfloat32float64zJInvalid input type. Expected `float32` or `float64`. Received: input type=zT`fft_length` must equal or larger than `sequence_length`. Received: sequence_length=z , fft_length=>ryhammingz]If a string is passed to `window`, it must be one of `"hann"`, `"hamming"`. Received: window=ryTperiodicrr rr r PThe shape of `window` must be equal to [sequence_length].Received: window shape=)rr FrR)n_fft hop_length win_lengthwindowcenterreturn_complex)rr TypeErrorrHrbstrrr hann_windowrhamming_windowrrrrTrrUstftswapaxesrfrg) r2rOrPrwrrwin need_unpackrYsamplesr[fft_unique_binsrZr&r&r'rs      rc Cst|}|jj}d} |durIt|tr-|dkr"tj|d|td} ntj|d|td} nt ||d} t | j dksA| j d|krIt d| j ||kr|dur| durd} |j ^} } } t |j d krud} t | }t||| | f}t|d d}tj||||| ||dd }| r|j d}t|g| |R}|St||}|||j d d}| dur||d }|||}t| ||ggd } |||}t| }| | }t|d|||fg}t|||g}tj|ddd}t||dg}t|||g}t| |d|} t|| }t||}|dur!dn|d }|dur/||}n |dur:|d  }n|}|d||fS)NryTr}rr r rFrSr)rrrrrlengthrrRconstantr)rD.)rirfrrbrrrrrrrrrHrTrrUristftrxrsquarer*tiledividemultiplyr`)r2rOrPrwrrrrhrrrrYrZrr[rexpected_output_lenl_padr_pad_sequence_lengthdenomoverlapsstartendr&r&r'rs                   rcCt|}t|Srj)rrrsqrtr2r&r&r'r rcCrrj)rrerfrr&r&r'rrrcCrrj)rrerfinvrr&r&r'rrrcCst|}t|}tj||Srj)rrrIsolve)abr&r&r'rr/rcCsLt|}t|jdkrt}nt|jt}t||}t j j ||||dS)Nr)ordr6rA) rrrrfloatxr result_typerrrrIru)r2rr7rDrr&r&r'rus   rucCrrj)rrlogdetrr&r&r'rrr)NF)T)rErj)ryT)NryT)NNF)$rTrkeras.src.backendrrkeras.src.backend.commonrkeras.src.backend.torch.corerrrkeras.src.backend.torch.numpyrr(r.r0r4r@rCrJrQr`rirkrprqrvrxrrrrrrrurr&r&r&r'sH        )    2   M c