| 145 | } |
| 146 | |
| 147 | TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { |
| 148 | TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); |
| 149 | TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); |
| 150 | |
| 151 | // Check type and shape of the input tensor |
| 152 | const TfLiteTensor* input = GetInput(context, node, kInputTensor); |
| 153 | TF_LITE_ENSURE(context, NumDimensions(input) >= 2); |
| 154 | if (input->type != kTfLiteFloat32) { |
| 155 | context->ReportError(context, |
| 156 | "Type '%s' for input is not supported by rfft2d.", |
| 157 | TfLiteTypeGetName(input->type)); |
| 158 | return kTfLiteError; |
| 159 | } |
| 160 | |
| 161 | // Check type and shape of the fft_length tensor |
| 162 | const TfLiteTensor* fft_length = GetInput(context, node, kFftLengthTensor); |
| 163 | const RuntimeShape fft_length_shape = GetTensorShape(fft_length); |
| 164 | |
| 165 | TF_LITE_ENSURE_EQ(context, NumDimensions(fft_length), 1); |
| 166 | TF_LITE_ENSURE_EQ(context, fft_length_shape.Dims(0), 2); |
| 167 | if (fft_length->type != kTfLiteInt32) { |
| 168 | context->ReportError(context, |
| 169 | "Type '%s' for fft_length is not supported by rfft2d.", |
| 170 | TfLiteTypeGetName(fft_length->type)); |
| 171 | return kTfLiteError; |
| 172 | } |
| 173 | |
| 174 | // Setup temporary tensors for fft computation. |
| 175 | TF_LITE_ENSURE_STATUS(InitTemporaryTensors(context, node)); |
| 176 | |
| 177 | // Set output type |
| 178 | TfLiteTensor* output = GetOutput(context, node, kOutputTensor); |
| 179 | output->type = kTfLiteComplex64; |
| 180 | |
| 181 | // Exit early if fft_length is a non-const tensor. Set output tensor and |
| 182 | // temporary tensors to dynamic, so that their tensor sizes can be determined |
| 183 | // in Eval. |
| 184 | if (!IsConstantTensor(fft_length)) { |
| 185 | TfLiteTensor* fft_integer_working_area = |
| 186 | GetTemporary(context, node, kFftIntegerWorkingAreaTensor); |
| 187 | TfLiteTensor* fft_double_working_area = |
| 188 | GetTemporary(context, node, kFftDoubleWorkingAreaTensor); |
| 189 | SetTensorToDynamic(fft_integer_working_area); |
| 190 | SetTensorToDynamic(fft_double_working_area); |
| 191 | SetTensorToDynamic(output); |
| 192 | return kTfLiteOk; |
| 193 | } |
| 194 | |
| 195 | TF_LITE_ENSURE_STATUS(ResizeOutputandTemporaryTensors(context, node)); |
| 196 | return kTfLiteOk; |
| 197 | } |
| 198 | |
| 199 | // Reorder the result so that it matches the pattern of tf.signal.rfft2d. |
| 200 | // In tf.signal.fft2d the frequency matrix of a 4x4 input is |
nothing calls this directly
no test coverage detected