license import { KernelConfig , KernelFunc , TensorInfo , Unpack , UnpackAttrs , UnpackInputs } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; import { reshape } from './Reshape' ; import { slice } from './Slice' ; export function unpack ( args ) { const { inputs , backend , attrs } = args ; const { value } = inputs ; let { axis } = attrs ; if ( axis < 0 ) { axis += value . shape . length ; } const valueRank = value . shape . length ; const num = value . shape [ axis ] ; const outShape = new Array ( valueRank - 1 ) ; let outIndex = 0 ; for ( let i = 0 ; i < valueRank ; i ++ ) { if ( i !== axis ) { outShape [ outIndex ++ ] = value . shape [ i ] ; } } const begin = new Array ( valueRank ) . fill ( 0 ) ; const size = value . shape . slice ( ) ; size [ axis ] = 1 ; const res = new Array ( num ) ; for ( let i = 0 ; i < res . length ; i ++ ) { begin [ axis ] = i ; const tempRes = slice ( { inputs : { x : value } , backend , attrs : { begin , size } } ) ; res [ i ] = reshape ( { inputs : { x : tempRes } , backend , attrs : { shape : outShape } } ) ; backend . disposeIntermediateTensorInfo ( tempRes ) ; } return res ; } export const unpackConfig = { kernelName : Unpack , backendName : 'cpu' , kernelFunc : unpack as { } as KernelFunc } ;