1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include <stdint.h>
17
18 #include <complex>
19 #include <vector>
20
21 #include <gtest/gtest.h>
22 #include "tensorflow/lite/interpreter.h"
23 #include "tensorflow/lite/kernels/custom_ops_register.h"
24 #include "tensorflow/lite/kernels/test_util.h"
25 #include "tensorflow/lite/schema/schema_generated.h"
26 #include "tensorflow/lite/testing/util.h"
27
28 namespace tflite {
29 namespace ops {
30 namespace custom {
31
32 namespace {
33
34 using ::testing::ElementsAreArray;
35
36 class Irfft2dOpModel : public SingleOpModel {
37 public:
Irfft2dOpModel(const TensorData & input,const TensorData & fft_lengths)38 Irfft2dOpModel(const TensorData& input, const TensorData& fft_lengths) {
39 input_ = AddInput(input);
40 fft_lengths_ = AddInput(fft_lengths);
41 TensorType output_type = TensorType_FLOAT32;
42 output_ = AddOutput({output_type, {}});
43
44 const std::vector<uint8_t> custom_option;
45 SetCustomOp("Irfft2d", custom_option, Register_IRFFT2D);
46 BuildInterpreter({GetShape(input_)});
47 }
48
input()49 int input() { return input_; }
fft_lengths()50 int fft_lengths() { return fft_lengths_; }
51
GetOutput()52 std::vector<float> GetOutput() { return ExtractVector<float>(output_); }
GetOutputShape()53 std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
54
55 private:
56 int input_;
57 int fft_lengths_;
58 int output_;
59 };
60
TEST(Irfft2dOpTest,FftLengthMatchesInputSize)61 TEST(Irfft2dOpTest, FftLengthMatchesInputSize) {
62 Irfft2dOpModel model({TensorType_COMPLEX64, {4, 3}}, {TensorType_INT32, {2}});
63 // clang-format off
64 model.PopulateTensor<std::complex<float>>(model.input(), {
65 {75, 0}, {-6, -1}, {9, 0}, {-10, 5}, {-3, 2}, {-6, 11},
66 {-15, 0}, {-2, 13}, {-5, 0}, {-10, -5}, {3, -6}, {-6, -11}
67 });
68 // clang-format on
69 model.PopulateTensor<int32_t>(model.fft_lengths(), {4, 4});
70 ASSERT_EQ(model.Invoke(), kTfLiteOk);
71
72 float expected_result[16] = {1, 2, 3, 4, 3, 8, 6, 3, 5, 2, 7, 6, 9, 5, 8, 3};
73 EXPECT_THAT(model.GetOutput(), ElementsAreArray(expected_result));
74 }
75
TEST(Irfft2dOpTest,FftLengthSmallerThanInputSize)76 TEST(Irfft2dOpTest, FftLengthSmallerThanInputSize) {
77 Irfft2dOpModel model({TensorType_COMPLEX64, {4, 3}}, {TensorType_INT32, {2}});
78 // clang-format off
79 model.PopulateTensor<std::complex<float>>(model.input(), {
80 {75, 0}, {-6, -1}, {9, 0}, {-10, 5}, {-3, 2}, {-6, 11},
81 {-15, 0}, {-2, 13}, {-5, 0}, {-10, -5}, {3, -6}, {-6, -11}
82 });
83 // clang-format on
84 model.PopulateTensor<int32_t>(model.fft_lengths(), {2, 2});
85 ASSERT_EQ(model.Invoke(), kTfLiteOk);
86
87 float expected_result[4] = {14, 18.5, 20.5, 22};
88 EXPECT_THAT(model.GetOutput(), ElementsAreArray(expected_result));
89 }
90
TEST(Irfft2dOpTest,FftLengthGreaterThanInputSize)91 TEST(Irfft2dOpTest, FftLengthGreaterThanInputSize) {
92 Irfft2dOpModel model({TensorType_COMPLEX64, {4, 3}}, {TensorType_INT32, {2}});
93 // clang-format off
94 model.PopulateTensor<std::complex<float>>(model.input(), {
95 {75, 0}, {-6, -1}, {9, 0}, {-10, 5}, {-3, 2}, {-6, 11},
96 {-15, 0}, {-2, 13}, {-5, 0}, {-10, -5}, {3, -6}, {-6, -11}
97 });
98 // clang-format on
99 model.PopulateTensor<int32_t>(model.fft_lengths(), {4, 8});
100 ASSERT_EQ(model.Invoke(), kTfLiteOk);
101
102 // clang-format off
103 float expected_result[32] = {
104 0.25, 0.54289322, 1.25, 1.25, 1.25, 1.95710678, 2.25, 1.25,
105 1.25, 2.85355339, 4.25, 3.91421356, 2.75, 2.14644661, 1.75, 1.08578644,
106 3., 1.43933983, 0.5, 2.14644661, 4., 3.56066017, 2.5, 2.85355339,
107 5.625, 3.65533009, 1.375, 3.3017767, 5.125, 2.59466991, 0.375, 2.9482233
108 };
109 // clang-format on
110 EXPECT_THAT(model.GetOutput(), ElementsAreArray(expected_result));
111 }
112
TEST(Irfft2dOpTest,InputDimsGreaterThan2)113 TEST(Irfft2dOpTest, InputDimsGreaterThan2) {
114 Irfft2dOpModel model({TensorType_COMPLEX64, {2, 2, 3}},
115 {TensorType_INT32, {2}});
116 // clang-format off
117 model.PopulateTensor<std::complex<float>>(model.input(), {
118 {30., 0.}, {-5, -3.}, { -4., 0.},
119 {-10., 0.}, {1., 7.}, { 0., 0.},
120 {58., 0.}, {-18., 6.}, { 26., 0.},
121 {-18., 0.}, { 14., 2.}, {-18., 0.}
122 });
123 // clang-format on
124 model.PopulateTensor<int32_t>(model.fft_lengths(), {2, 4});
125 ASSERT_EQ(model.Invoke(), kTfLiteOk);
126
127 float expected_result[16] = {1., 2., 3., 4., 3., 8., 6., 3.,
128 5., 2., 7., 6., 7., 3., 23., 5.};
129 EXPECT_THAT(model.GetOutput(), ElementsAreArray(expected_result));
130 }
131
132 } // namespace
133 } // namespace custom
134 } // namespace ops
135 } // namespace tflite
136