license import { KernelConfig , KernelFunc , Pack , PackAttrs , PackInputs , TensorInfo , util } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { concat } from './Concat' ; import { expandDims } from './ExpandDims' ; export function pack ( args ) { const { inputs , backend , attrs } = args ; const { axis } = attrs ; if ( inputs . length === 1 ) { return expandDims ( { inputs : { input : inputs [ 0 ] } , backend , attrs : { dim : axis } } ) ; } const shape = inputs [ 0 ] . shape ; const dtype = inputs [ 0 ] . dtype ; inputs . forEach ( t => { util . assertShapesMatch ( shape , t . shape , 'All tensors passed to stack must have matching shapes' ) ; util . assert ( dtype === t . dtype , ( ) => 'All tensors passed to stack must have matching dtypes' ) ; } ) ; const intermediateTensorInfos = [ ] ; const expandedTensors = inputs . map ( t => { const expandedT = expandDims ( { inputs : { input : t } , backend , attrs : { dim : axis } } ) ; intermediateTensorInfos . push ( expandedT ) ; return expandedT ; } ) ; const result = concat ( { inputs : expandedTensors , backend , attrs : { axis } } ) ; intermediateTensorInfos . forEach ( t => backend . disposeIntermediateTensorInfo ( t ) ) ; return result ; } export const packConfig = { kernelName : Pack , backendName : 'webgl' , kernelFunc : pack as { } as KernelFunc } ;