license import { KernelConfig , KernelFunc , Select , SelectInputs , TensorInfo , TypedArray , upcastType , util } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; import { assertNotComplex } from '../cpu_util' ; export function select ( args ) { const { inputs , backend } = args ; const { condition , t , e } = inputs ; assertNotComplex ( [ condition , t , e ] , 'select' ) ; const conditionRank = condition . shape . length ; const values = backend . data . get ( condition . dataId ) . values as TypedArray ; const tValues = backend . data . get ( t . dataId ) . values as TypedArray ; const eValues = backend . data . get ( e . dataId ) . values as TypedArray ; const resultDtype = upcastType ( t . dtype , e . dtype ) ; const newValues = util . makeZerosTypedArray ( util . sizeFromShape ( t . shape ) , resultDtype ) ; let index = 0 ; const offset = conditionRank === 0 || conditionRank > 1 || t . shape . length === 1 ? 1 : util . sizeFromShape ( t . shape . slice ( 1 ) ) ; for ( let i = 0 ; i < values . length ; i ++ ) { for ( let j = 0 ; j < offset ; j ++ ) { if ( values [ i ] === 1 ) { newValues [ index ++ ] = tValues [ i ] ; } else { newValues [ index ++ ] = eValues [ i ] ; } } } return backend . makeTensorInfo ( t . shape , resultDtype , newValues ) ; } export const selectConfig = { kernelName : Select , backendName : 'cpu' , kernelFunc : select as { } as KernelFunc } ;