license import { FFT , FFTInputs , KernelConfig , KernelFunc , TensorInfo , util } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; import { fftBatch } from '../utils/fft_utils' ; import { reshape } from './Reshape' ; export function fft ( args ) { const { inputs , backend } = args ; const { input } = inputs ; const inputSize = util . sizeFromShape ( input . shape ) ; const innerDimensionSize = input . shape [ input . shape . length - 1 ] ; const batch = inputSize / innerDimensionSize ; const input2D = reshape ( { inputs : { x : input } , backend , attrs : { shape : [ batch , innerDimensionSize ] } } ) ; const result = fftBatch ( input2D , false , backend ) ; const resultReshaped = reshape ( { inputs : { x : result } , backend , attrs : { shape : input . shape } } ) ; backend . disposeIntermediateTensorInfo ( input2D ) ; backend . disposeIntermediateTensorInfo ( result ) ; return resultReshaped ; } export const fftConfig = { kernelName : FFT , backendName : 'cpu' , kernelFunc : fft as { } as KernelFunc } ;