Function fftImpl

tfjs-backend-cpu/src/utils/fft_utils.ts:97–163

(
    input: TensorInfo, inverse: boolean,
    cpuBackend: MathBackendCPU)

}

export function fftImpl(
    input: TensorInfo, inverse: boolean,
    cpuBackend: MathBackendCPU): {real: Float32Array, imag: Float32Array} {
  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: TensorInfo =
          cpuBackend.makeTensorInfo(resultShape, 'float32', result.real);
      const imagInfo: TensorInfo =
          cpuBackend.makeTensorInfo(resultShape, 'float32', result.imag);

      const sizeInfo: TensorInfo = cpuBackend.makeTensorInfo(
          [], 'float32',
          util.createScalarValue(inputSize as unknown 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;