license import { ArgMin , ArgMinAttrs , ArgMinInputs , backend_util , KernelConfig , KernelFunc , TensorInfo , util } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { argMinMaxReduce } from '../kernel_utils/arg_min_max' ; import { transpose } from './Transpose' ; export function argMin ( args ) { const { inputs , backend , attrs } = args ; const { x } = inputs ; const { axis } = attrs ; let axes = util . parseAxisParam ( axis , x . shape ) ; const permutedAxes = backend_util . getAxesPermutation ( axes , x . shape . length ) ; let $x = x ; const intermediateTensorInfos = [ ] ; if ( permutedAxes != null ) { $x = transpose ( { inputs : { x } , backend , attrs : { perm : permutedAxes } } ) ; intermediateTensorInfos . push ( $x ) ; axes = backend_util . getInnerMostAxes ( axes . length , $x . shape . length ) ; } backend_util . assertAxesAreInnerMostDims ( 'argMin' , [ axes [ 0 ] ] , $x . shape . length ) ; const out = argMinMaxReduce ( backend , $x , axes [ 0 ] , 'min' ) ; intermediateTensorInfos . forEach ( t => backend . disposeIntermediateTensorInfo ( t ) ) ; return out ; } export const argMinConfig = { kernelName : ArgMin , backendName : 'webgl' , kernelFunc : argMin as { } as KernelFunc } ;