license import { backend_util , Tensor , TensorInfo , TypedArray , util } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; import { add } from '../kernels/Add' ; import { complex } from '../kernels/Complex' ; import { concat } from '../kernels/Concat' ; import { identity } from '../kernels/Identity' ; import { imag } from '../kernels/Imag' ; import { multiply } from '../kernels/Multiply' ; import { real } from '../kernels/Real' ; import { realDivConfig } from '../kernels/RealDiv' ; import { slice } from '../kernels/Slice' ; import { sub } from '../kernels/Sub' ; /**
 * Calculate FFT of inner most elements of batch tensor.
 */ export function fftBatch ( input , inverse , cpuBackend ) { const inputShape = input . shape ; const batch = inputShape [ 0 ] ; const innerDim = inputShape [ 1 ] ; const inputVals = cpuBackend . data . get ( input . dataId ) ; const real2D = inputVals . complexTensorInfos . real ; const imag2D = inputVals . complexTensorInfos . imag ; const resultShape = [ batch , innerDim ] ; const resultSize = util . sizeFromShape ( resultShape ) ; const resultReal = util . getTypedArrayFromDType ( 'float32' , resultSize ) ; const resultImag = util . getTypedArrayFromDType ( 'float32' , resultSize ) ; for ( let b = 0 ; b < batch ; b ++ ) { const r = slice ( { inputs : { x : real2D } , backend : cpuBackend , attrs : { begin : [ b , 0 ] , size : [ 1 , innerDim ] } } ) ; const i = slice ( { inputs : { x : imag2D } , backend : cpuBackend , attrs : { begin : [ b , 0 ] , size : [ 1 , innerDim ] } } ) ; const input = complex ( { inputs : { real : r , imag : i } , backend : cpuBackend } ) ; const { real , imag } = fftImpl ( input , inverse , cpuBackend ) ; const res = backend_util . mergeRealAndImagArrays ( real , imag ) ; for ( let d = 0 ; d < innerDim ; d ++ ) { const c = backend_util . getComplexWithIndex ( res , d ) ; resultReal [ b * innerDim + d ] = c . real ; resultImag [ b * innerDim + d ] = c . imag ; } cpuBackend . disposeIntermediateTensorInfo ( r ) ; cpuBackend . disposeIntermediateTensorInfo ( i ) ; cpuBackend . disposeIntermediateTensorInfo ( input ) ; } const $realInfo = cpuBackend . makeTensorInfo ( resultShape , 'float32' , resultReal ) ; const $imagInfo = cpuBackend . makeTensorInfo ( resultShape , 'float32' , resultImag ) ; const result = complex ( { inputs : { real : $realInfo , imag : $imagInfo } , backend : cpuBackend } ) ; cpuBackend . disposeIntermediateTensorInfo ( $realInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( $imagInfo ) ; return result ; } export function fftImpl ( input , inverse , cpuBackend ) { const inputSize = util . sizeFromShape ( input . shape ) ; const inputVals = cpuBackend . data . get ( input . dataId ) ; const realVals = cpuBackend . data . get ( inputVals . complexTensorInfos . real . dataId ) . values as Float32Array ; const imagVals = cpuBackend . data . get ( inputVals . complexTensorInfos . imag . dataId ) . values as Float32Array ; if ( isExponentOf2 ( inputSize ) ) { const result = fftRadix2 ( realVals , imagVals , inputSize , inverse , cpuBackend ) ; const resultShape = [ input . shape [ 0 ] , input . shape [ 1 ] ] ; if ( inverse ) { const realInfo = cpuBackend . makeTensorInfo ( resultShape , 'float32' , result . real ) ; const imagInfo = cpuBackend . makeTensorInfo ( resultShape , 'float32' , result . imag ) ; const sizeInfo = cpuBackend . makeTensorInfo ( [ ] , 'float32' , util . createScalarValue ( inputSize as { } as 'float32' , 'float32' ) ) ; const sizeInfoCopy = identity ( { inputs : { x : sizeInfo } , backend : cpuBackend } ) ; const divRealInfo = realDivConfig . kernelFunc ( { inputs : { a : realInfo , b : sizeInfo } , backend : cpuBackend } ) as TensorInfo ; const divImagInfo = realDivConfig . kernelFunc ( { inputs : { a : imagInfo , b : sizeInfoCopy } , backend : cpuBackend } ) as TensorInfo ; const divRealVals = cpuBackend . data . get ( divRealInfo . dataId ) . values as Float32Array ; const divImagVals = cpuBackend . data . get ( divImagInfo . dataId ) . values as Float32Array ; cpuBackend . disposeIntermediateTensorInfo ( realInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( imagInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( sizeInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( sizeInfoCopy ) ; cpuBackend . disposeIntermediateTensorInfo ( divRealInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( divImagInfo ) ; return { real : divRealVals , imag : divImagVals } ; } return result ; } else { const data = backend_util . mergeRealAndImagArrays ( realVals , imagVals ) ; const rawOutput = fourierTransformByMatmul ( data , inputSize , inverse ) as Float32Array ; return backend_util . splitRealAndImagArrays ( rawOutput ) ; } } function isExponentOf2 ( size ) { return ( size & size - 1 ) === 0 ; } function fftRadix2 ( realVals , imagVals , size , inverse , cpuBackend ) { if ( size === 1 ) { return { real : realVals , imag : imagVals } ; } const data = backend_util . mergeRealAndImagArrays ( realVals , imagVals ) ; const half = size / 2 ; const evenComplex = backend_util . complexWithEvenIndex ( data ) ; const evenRealVals = evenComplex . real ; const evenImagVals = evenComplex . imag ; const evenShape = [ evenRealVals . length ] ; const evenRealInfo = cpuBackend . makeTensorInfo ( evenShape , 'float32' , evenRealVals ) ; const evenImagInfo = cpuBackend . makeTensorInfo ( evenShape , 'float32' , evenImagVals ) ; const evenTensorInfo = complex ( { inputs : { real : evenRealInfo , imag : evenImagInfo } , backend : cpuBackend } ) ; const oddComplex = backend_util . complexWithOddIndex ( data ) ; const oddRealVals = oddComplex . real ; const oddImagVals = oddComplex . imag ; const oddShape = [ oddRealVals . length ] ; const oddRealInfo = cpuBackend . makeTensorInfo ( oddShape , 'float32' , oddRealVals ) ; const oddImagInfo = cpuBackend . makeTensorInfo ( oddShape , 'float32' , oddImagVals ) ; const oddTensorInfo = complex ( { inputs : { real : oddRealInfo , imag : oddImagInfo } , backend : cpuBackend } ) ; const $evenComplex = fftRadix2 ( evenRealVals , evenImagVals , half , inverse , cpuBackend ) ; const $evenRealVals = $evenComplex . real ; const $evenImagVals = $evenComplex . imag ; const $evenShape = [ $evenRealVals . length ] ; const $evenRealInfo = cpuBackend . makeTensorInfo ( $evenShape , 'float32' , $evenRealVals ) ; const $evenImagInfo = cpuBackend . makeTensorInfo ( $evenShape , 'float32' , $evenImagVals ) ; const $evenTensorInfo = complex ( { inputs : { real : $evenRealInfo , imag : $evenImagInfo } , backend : cpuBackend } ) ; const $oddComplex = fftRadix2 ( oddRealVals , oddImagVals , half , inverse , cpuBackend ) ; const $oddRealVals = $oddComplex . real ; const $oddImagVals = $oddComplex . imag ; const $oddShape = [ $oddRealVals . length ] ; const $oddRealInfo = cpuBackend . makeTensorInfo ( $oddShape , 'float32' , $oddRealVals ) ; const $oddImagInfo = cpuBackend . makeTensorInfo ( $oddShape , 'float32' , $oddImagVals ) ; const $oddTensorInfo = complex ( { inputs : { real : $oddRealInfo , imag : $oddImagInfo } , backend : cpuBackend } ) ; const e = backend_util . exponents ( size , inverse ) ; const eShape = [ e . real . length ] ; const eRealInfo = cpuBackend . makeTensorInfo ( eShape , 'float32' , e . real ) ; const eImagInfo = cpuBackend . makeTensorInfo ( eShape , 'float32' , e . imag ) ; const complexInfo = complex ( { inputs : { real : eRealInfo , imag : eImagInfo } , backend : cpuBackend } ) ; const exponentInfo = multiply ( { inputs : { a : complexInfo , b : $oddTensorInfo } , backend : cpuBackend } ) as TensorInfo ; const addPart = add ( { inputs : { a : $evenTensorInfo , b : exponentInfo } , backend : cpuBackend } ) as TensorInfo ; const subPart = sub ( { inputs : { a : $evenTensorInfo , b : exponentInfo } , backend : cpuBackend } ) as TensorInfo ; const addPartReal = real ( { inputs : { input : addPart } , backend : cpuBackend } ) ; const subPartReal = real ( { inputs : { input : subPart } , backend : cpuBackend } ) ; const addPartImag = imag ( { inputs : { input : addPart } , backend : cpuBackend } ) ; const subPartImag = imag ( { inputs : { input : subPart } , backend : cpuBackend } ) ; const $real = concat ( { inputs : [ addPartReal as Tensor , subPartReal as Tensor ] , backend : cpuBackend , attrs : { axis : 0 } } ) ; const $imag = concat ( { inputs : [ addPartImag as Tensor , subPartImag as Tensor ] , backend : cpuBackend , attrs : { axis : 0 } } ) ; const $realVals = cpuBackend . data . get ( $real . dataId ) . values as Float32Array ; const $imagVals = cpuBackend . data . get ( $imag . dataId ) . values as Float32Array ; cpuBackend . disposeIntermediateTensorInfo ( evenRealInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( evenImagInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( evenTensorInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( oddRealInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( oddImagInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( oddTensorInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( $evenRealInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( $evenImagInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( $evenTensorInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( $oddRealInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( $oddImagInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( $oddTensorInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( eRealInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( eImagInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( complexInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( exponentInfo ) ; cpuBackend . disposeIntermediateTensorInfo ( addPart ) ; cpuBackend . disposeIntermediateTensorInfo ( subPart ) ; cpuBackend . disposeIntermediateTensorInfo ( addPartReal ) ; cpuBackend . disposeIntermediateTensorInfo ( addPartImag ) ; cpuBackend . disposeIntermediateTensorInfo ( subPartReal ) ; cpuBackend . disposeIntermediateTensorInfo ( subPartImag ) ; cpuBackend . disposeIntermediateTensorInfo ( $real ) ; cpuBackend . disposeIntermediateTensorInfo ( $imag ) ; return { real : $realVals , imag : $imagVals } ; } function fourierTransformByMatmul ( data , size , inverse ) { const ret = new Float32Array ( size * 2 ) ; for ( let r = 0 ; r < size ; r ++ ) { let real = 0.0 ; let imag = 0.0 ; for ( let c = 0 ; c < size ; c ++ ) { const e = backend_util . exponent ( r * c , size , inverse ) ; const term = backend_util . getComplexWithIndex ( data as Float32Array , c ) ; real += term . real * e . real - term . imag * e . imag ; imag += term . real * e . imag + term . imag * e . real ; } if ( inverse ) { real /= size ; imag /= size ; } backend_util . assignToTypedArray ( ret , real , imag , r ) ; } return ret ; }