license import { KernelConfig , KernelFunc , OneHot , OneHotAttrs , OneHotInputs , TensorInfo , TypedArray , util } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; import { assertNotComplex } from '../cpu_util' ; export function oneHot ( args ) { const { inputs , backend , attrs } = args ; const { indices } = inputs ; const { depth , onValue , offValue } = attrs ; assertNotComplex ( indices , 'oneHot' ) ; const indicesSize = util . sizeFromShape ( indices . shape ) ; const res = new Float32Array ( indicesSize * depth ) ; res . fill ( offValue ) ; const indicesVal = backend . data . get ( indices . dataId ) . values as TypedArray ; for ( let event = 0 ; event < indicesSize ; ++ event ) { if ( indicesVal [ event ] >= 0 && indicesVal [ event ] < depth ) { res [ event * depth + indicesVal [ event ] ] = onValue ; } } return backend . makeTensorInfo ( [ ... indices . shape , depth ] , 'int32' , res ) ; } export const oneHotConfig = { kernelName : OneHot , backendName : 'cpu' , kernelFunc : oneHot as { } as KernelFunc } ;