license import { backend_util , ConcatInputs , env , TensorInfo , util } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { ConcatProgram } from '../concat_gpu' ; import { ConcatPackedProgram } from '../concat_packed_gpu' ; import { concatImplCPU } from '../kernel_utils/shared' ; import { complex } from './Complex' ; import { imag } from './Imag' ; import { real } from './Real' ; import { reshape } from './Reshape' ; export function concatImpl ( inputs , axis , backend ) { const dtype = inputs [ 0 ] . dtype ; if ( dtype === 'complex64' ) { const reals = inputs . map ( ( t ) => real ( { inputs : { input : t } , backend } ) ) ; const imags = inputs . map ( ( t ) => imag ( { inputs : { input : t } , backend } ) ) ; const realConcated = concatImpl ( reals , axis , backend ) ; const imagConcated = concatImpl ( imags , axis , backend ) ; const result = complex ( { inputs : { real : realConcated , imag : imagConcated } , backend } ) ; reals . forEach ( r => backend . disposeIntermediateTensorInfo ( r ) ) ; imags . forEach ( i => backend . disposeIntermediateTensorInfo ( i ) ) ; backend . disposeIntermediateTensorInfo ( realConcated ) ; backend . disposeIntermediateTensorInfo ( imagConcated ) ; return result ; } let runOnCpu = backend . shouldExecuteOnCPU ( inputs ) ; if ( dtype === 'string' ) { runOnCpu = true ; } if ( runOnCpu ) { const tensors2D = inputs . map ( t => { const innerSize = util . sizeFromShape ( t . shape . slice ( axis ) ) ; const shape = [ - 1 , innerSize ] ; return reshape ( { inputs : { x : t } , backend , attrs : { shape } } ) ; } ) ; const inputsValShapes = tensors2D . map ( t => { return { vals : backend . readSync ( t . dataId ) , shape : t . shape } ; } ) ; const outShape = backend_util . computeOutShape ( tensors2D . map ( t => t . shape ) , 1 ) ; const simplyConcat = tensors2D [ 0 ] . shape [ 0 ] === 1 ; const outVals = concatImplCPU ( inputsValShapes , outShape , dtype , simplyConcat ) ; const finalOutShape = backend_util . computeOutShape ( inputs . map ( t => t . shape ) , axis ) ; const outInfo = backend . makeTensorInfo ( finalOutShape , dtype , outVals ) ; tensors2D . forEach ( t => backend . disposeIntermediateTensorInfo ( t ) ) ; return outInfo ; } if ( inputs . length > env ( ) . getNumber ( 'WEBGL_MAX_TEXTURES_IN_SHADER' ) ) { const midIndex = Math . floor ( inputs . length / 2 ) ; const leftSide = concatImpl ( inputs . slice ( 0 , midIndex ) , axis , backend ) ; const rightSide = concatImpl ( inputs . slice ( midIndex ) , axis , backend ) ; const result = concatImpl ( [ leftSide , rightSide ] , axis , backend ) ; backend . disposeIntermediateTensorInfo ( leftSide ) ; backend . disposeIntermediateTensorInfo ( rightSide ) ; return result ; } if ( env ( ) . getBool ( 'WEBGL_PACK_ARRAY_OPERATIONS' ) && inputs [ 0 ] . shape . length > 1 ) { const program = new ConcatPackedProgram ( inputs . map ( t => t . shape ) , axis ) ; return backend . runWebGLProgram ( program , inputs , dtype ) ; } const { tensors2D , outShape } = computeTensors2D ( inputs , axis , backend ) ; const program = new ConcatProgram ( tensors2D . map ( t => t . shape as [ number , number ] ) ) ; const result = backend . runWebGLProgram ( program , tensors2D , dtype ) ; tensors2D . forEach ( r => backend . disposeIntermediateTensorInfo ( r ) ) ; const reshapedResult = reshape ( { inputs : { x : result } , attrs : { shape : outShape } , backend } ) ; backend . disposeIntermediateTensorInfo ( result ) ; return reshapedResult ; } function computeTensors2D ( inputs , axis , backend ) { const outShape = backend_util . computeOutShape ( inputs . map ( t => t . shape ) , axis ) ; const tensors2D = inputs . map ( x => reshape ( { inputs : { x } , attrs : { shape : [ - 1 , util . sizeFromShape ( x . shape . slice ( axis ) ) ] } , backend } ) ) ; return { tensors2D , outShape } ; }