license import { TensorInfo } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { ReshapePackedProgram } from '../reshape_packed_gpu' ; import { getBatchDim , getRowsCols } from '../webgl_util' ; export function packedReshape ( input , afterShape , backend ) { const input3DShape = [ getBatchDim ( input . shape ) , ... getRowsCols ( input . shape ) ] as [ number , number , number ] ; const input3D = { dtype : input . dtype , shape : input3DShape , dataId : input . dataId } ; const afterShapeAs3D = [ getBatchDim ( afterShape ) , ... getRowsCols ( afterShape ) ] as [ number , number , number ] ; const program = new ReshapePackedProgram ( afterShapeAs3D , input3DShape ) ; const preventEagerUnpackingOfOutput = true ; const customValues = [ input3DShape ] ; const output = backend . runWebGLProgram ( program , [ input3D ] , input . dtype , customValues , preventEagerUnpackingOfOutput ) ; return { dataId : output . dataId , shape : afterShape , dtype : output . dtype } ; }