license import { KernelConfig , KernelFunc , OneHot , OneHotAttrs , OneHotInputs , TensorInfo , util } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { OneHotProgram } from '../onehot_gpu' ; import { reshape } from './Reshape' ; export const oneHot = ( args ) : TensorInfo => { const { inputs , backend , attrs } = args ; const { indices } = inputs ; const { depth , onValue , offValue } = attrs ; const indicesSize = util . sizeFromShape ( indices . shape ) ; const program = new OneHotProgram ( indicesSize , depth , onValue , offValue ) ; const reshaped = reshape ( { inputs : { x : indices } , backend , attrs : { shape : [ indicesSize ] } } ) ; const result = backend . runWebGLProgram ( program , [ reshaped ] , indices . dtype ) ; backend . disposeIntermediateTensorInfo ( reshaped ) ; const outShape = [ ... indices . shape , depth ] ; const out = reshape ( { inputs : { x : result } , backend , attrs : { shape : outShape } } ) ; backend . disposeIntermediateTensorInfo ( result ) ; return out ; } ; export const oneHotConfig = { kernelName : OneHot , backendName : 'webgl' , kernelFunc : oneHot as { } as KernelFunc } ;