license import { TensorInfo , util } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { FFTProgram } from '../fft_gpu' ; import { complex } from './Complex' ; import { reshape } from './Reshape' ; export function fftImpl ( x , inverse , backend ) { const xData = backend . texData . get ( x . dataId ) ; const inputSize = util . sizeFromShape ( x . shape ) ; const innerDimensionSize = x . shape [ x . shape . length - 1 ] ; const batch = inputSize / innerDimensionSize ; const input2D = reshape ( { inputs : { x } , backend , attrs : { shape : [ batch , innerDimensionSize ] } } ) ; const xShape = input2D . shape as [ number , number ] ; const realProgram = new FFTProgram ( 'real' , xShape , inverse ) ; const imagProgram = new FFTProgram ( 'imag' , xShape , inverse ) ; const inputs = [ { dataId : xData . complexTensorInfos . real . dataId , dtype : xData . complexTensorInfos . real . dtype , shape : xShape } , { dataId : xData . complexTensorInfos . imag . dataId , dtype : xData . complexTensorInfos . imag . dtype , shape : xShape } ] ; const realPart = backend . runWebGLProgram ( realProgram , inputs , 'float32' ) ; const imagPart = backend . runWebGLProgram ( imagProgram , inputs , 'float32' ) ; const complexOutput = complex ( { inputs : { real : realPart , imag : imagPart } , backend } ) ; backend . disposeIntermediateTensorInfo ( realPart ) ; backend . disposeIntermediateTensorInfo ( imagPart ) ; const complexOutputReshaped = reshape ( { inputs : { x : complexOutput } , backend , attrs : { shape : x . shape } } ) ; backend . disposeIntermediateTensorInfo ( input2D ) ; backend . disposeIntermediateTensorInfo ( complexOutput ) ; return complexOutputReshaped ; }