Function RFFTShape

tensorflow/core/ops/spectral_ops.cc:75–109

    });

Status RFFTShape(InferenceContext* c, const bool forward, const int rank) {
  ShapeHandle out;
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), rank, &out));

  // Check that fft_length has shape [rank].
  ShapeHandle unused_shape;
  DimensionHandle unused_dim;
  ShapeHandle fft_length_input = c->input(1);
  TF_RETURN_IF_ERROR(c->WithRank(fft_length_input, 1, &unused_shape));
  TF_RETURN_IF_ERROR(
      c->WithValue(c->Dim(fft_length_input, 0), rank, &unused_dim));
  const Tensor* fft_length_tensor = c->input_tensor(1);

  // If fft_length is unknown at graph creation time, we can't predict the
  // output size.
  if (fft_length_tensor == nullptr) {
    // We can't know the dimension of any of the rank inner dimensions of the
    // output without knowing fft_length.
    for (int i = 0; i < rank; ++i) {
      TF_RETURN_IF_ERROR(c->ReplaceDim(out, -rank + i, c->UnknownDim(), &out));
    }
  } else {
    auto fft_length_as_vec = fft_length_tensor->vec<int32>();
    for (int i = 0; i < rank; ++i) {
      // For RFFT, replace the last dimension with fft_length/2 + 1.
      auto dim = forward && i == rank - 1 && fft_length_as_vec(i) != 0
                     ? fft_length_as_vec(i) / 2 + 1
                     : fft_length_as_vec(i);
      TF_RETURN_IF_ERROR(c->ReplaceDim(out, -rank + i, c->MakeDim(dim), &out));
    }
  }

  c->set_output(0, out);
  return Status::OK();
}

REGISTER_OP("RFFT")
    .Input("input: float")