Method DoFFT

tensorflow/core/kernels/fft_ops.cc:126–251

  bool IsReal() const override { return _Real; }

  void DoFFT(OpKernelContext* ctx, const Tensor& in, uint64* fft_shape,
             Tensor* out) override {
    // Create the axes (which are always trailing).
    const auto axes = Eigen::ArrayXi::LinSpaced(FFTRank, 1, FFTRank);
    auto device = ctx->eigen_device<CPUDevice>();

    if (!IsReal()) {
      // Compute the FFT using Eigen.
      constexpr auto direction =
          Forward ? Eigen::FFT_FORWARD : Eigen::FFT_REVERSE;
      if (in.dtype() == DT_COMPLEX64) {
        DCHECK_EQ(out->dtype(), DT_COMPLEX64);
        auto input = Tensor(in).flat_inner_dims<complex64, FFTRank + 1>();
        auto output = out->flat_inner_dims<complex64, FFTRank + 1>();
        output.device(device) =
            input.template fft<Eigen::BothParts, direction>(axes);
      } else {
        DCHECK_EQ(DT_COMPLEX128, in.dtype());
        DCHECK_EQ(DT_COMPLEX128, out->dtype());
        auto input = Tensor(in).flat_inner_dims<complex128, FFTRank + 1>();
        auto output = out->flat_inner_dims<complex128, FFTRank + 1>();
        output.device(device) =
            input.template fft<Eigen::BothParts, direction>(axes);
      }
    } else {
      if (IsForward()) {
        auto input = Tensor(in).flat_inner_dims<float, FFTRank + 1>();
        const auto input_dims = input.dimensions();

        // Slice input to fft_shape on its inner-most dimensions.
        Eigen::DSizes<Eigen::DenseIndex, FFTRank + 1> input_slice_sizes;
        input_slice_sizes[0] = input_dims[0];
        TensorShape temp_shape{input_dims[0]};
        for (int i = 1; i <= FFTRank; ++i) {
          input_slice_sizes[i] = fft_shape[i - 1];
          temp_shape.AddDim(fft_shape[i - 1]);
        }

        auto output = out->flat_inner_dims<complex64, FFTRank + 1>();
        const Eigen::DSizes<Eigen::DenseIndex, FFTRank + 1> zero_start_indices;

        // Compute the full FFT using a temporary tensor.
        Tensor temp;
        OP_REQUIRES_OK(ctx, ctx->allocate_temp(DataTypeToEnum<complex64>::v(),
                                               temp_shape, &temp));
        auto full_fft = temp.flat_inner_dims<complex64, FFTRank + 1>();
        full_fft.device(device) =
            input.slice(zero_start_indices, input_slice_sizes)
                .template fft<Eigen::BothParts, Eigen::FFT_FORWARD>(axes);

        // Slice away the negative frequency components.
        output.device(device) =
            full_fft.slice(zero_start_indices, output.dimensions());
      } else {
        // Reconstruct the full FFT and take the inverse.
        auto input = Tensor(in).flat_inner_dims<complex64, FFTRank + 1>();
        auto output = out->flat_inner_dims<float, FFTRank + 1>();
        const auto input_dims = input.dimensions();