1 /* Copyright 2017 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 "tensorflow/core/framework/common_shape_fns.h" 17 #include "tensorflow/core/framework/numeric_op.h" 18 #include "tensorflow/core/framework/op.h" 19 #include "tensorflow/core/framework/shape_inference.h" 20 21 namespace tensorflow { 22 23 using shape_inference::DimensionHandle; 24 using shape_inference::InferenceContext; 25 using shape_inference::ShapeHandle; 26 27 REGISTER_OP("FFT") 28 .Input("input: complex64") 29 .Output("output: complex64") 30 .SetShapeFn([](InferenceContext* c) { 31 return shape_inference::UnchangedShapeWithRankAtLeast(c, 1); 32 }); 33 34 REGISTER_OP("IFFT") 35 .Input("input: complex64") 36 .Output("output: complex64") 37 .SetShapeFn([](InferenceContext* c) { 38 return shape_inference::UnchangedShapeWithRankAtLeast(c, 1); 39 }); 40 41 REGISTER_OP("FFT2D") 42 .Input("input: complex64") 43 .Output("output: complex64") 44 .SetShapeFn([](InferenceContext* c) { 45 return shape_inference::UnchangedShapeWithRankAtLeast(c, 2); 46 }); 47 48 REGISTER_OP("IFFT2D") 49 .Input("input: complex64") 50 .Output("output: complex64") 51 .SetShapeFn([](InferenceContext* c) { 52 return shape_inference::UnchangedShapeWithRankAtLeast(c, 2); 53 }); 54 55 REGISTER_OP("FFT3D") 56 .Input("input: complex64") 57 .Output("output: complex64") 58 .SetShapeFn([](InferenceContext* c) { 59 return shape_inference::UnchangedShapeWithRankAtLeast(c, 3); 60 }); 61 62 REGISTER_OP("IFFT3D") 63 .Input("input: complex64") 64 .Output("output: complex64") 65 .SetShapeFn([](InferenceContext* c) { 66 return shape_inference::UnchangedShapeWithRankAtLeast(c, 3); 67 }); 68 69 Status RFFTShape(InferenceContext* c, const bool forward, const int rank) { 70 ShapeHandle out; 71 TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), rank, &out)); 72 73 // Check that fft_length has shape [rank]. 74 ShapeHandle unused_shape; 75 DimensionHandle unused_dim; 76 ShapeHandle fft_length_input = c->input(1); 77 TF_RETURN_IF_ERROR(c->WithRank(fft_length_input, 1, &unused_shape)); 78 TF_RETURN_IF_ERROR( 79 c->WithValue(c->Dim(fft_length_input, 0), rank, &unused_dim)); 80 const Tensor* fft_length_tensor = c->input_tensor(1); 81 82 // If fft_length is unknown at graph creation time, we can't predict the 83 // output size. 84 if (fft_length_tensor == nullptr) { 85 // We can't know the dimension of any of the rank inner dimensions of the 86 // output without knowing fft_length. 87 for (int i = 0; i < rank; ++i) { 88 TF_RETURN_IF_ERROR(c->ReplaceDim(out, -rank + i, c->UnknownDim(), &out)); 89 } 90 } else { 91 auto fft_length_as_vec = fft_length_tensor->vec<int32>(); 92 for (int i = 0; i < rank; ++i) { 93 // For RFFT, replace the last dimension with fft_length/2 + 1. 94 auto dim = forward && i == rank - 1 && fft_length_as_vec(i) != 0 95 ? fft_length_as_vec(i) / 2 + 1 96 : fft_length_as_vec(i); 97 TF_RETURN_IF_ERROR(c->ReplaceDim(out, -rank + i, c->MakeDim(dim), &out)); 98 } 99 } 100 101 c->set_output(0, out); 102 return Status::OK(); 103 } 104 105 REGISTER_OP("RFFT") 106 .Input("input: float") 107 .Input("fft_length: int32") 108 .Output("output: complex64") 109 .SetShapeFn([](InferenceContext* c) { return RFFTShape(c, true, 1); }); 110 111 REGISTER_OP("IRFFT") 112 .Input("input: complex64") 113 .Input("fft_length: int32") 114 .Output("output: float") 115 .SetShapeFn([](InferenceContext* c) { return RFFTShape(c, false, 1); }); 116 117 REGISTER_OP("RFFT2D") 118 .Input("input: float") 119 .Input("fft_length: int32") 120 .Output("output: complex64") 121 .SetShapeFn([](InferenceContext* c) { return RFFTShape(c, true, 2); }); 122 123 REGISTER_OP("IRFFT2D") 124 .Input("input: complex64") 125 .Input("fft_length: int32") 126 .Output("output: float") 127 .SetShapeFn([](InferenceContext* c) { return RFFTShape(c, false, 2); }); 128 129 REGISTER_OP("RFFT3D") 130 .Input("input: float") 131 .Input("fft_length: int32") 132 .Output("output: complex64") 133 .SetShapeFn([](InferenceContext* c) { return RFFTShape(c, true, 3); }); 134 135 REGISTER_OP("IRFFT3D") 136 .Input("input: complex64") 137 .Input("fft_length: int32") 138 .Output("output: float") 139 .SetShapeFn([](InferenceContext* c) { return RFFTShape(c, false, 3); }); 140 141 // Deprecated ops: 142 REGISTER_OP("BatchFFT") 143 .Input("input: complex64") 144 .Output("output: complex64") 145 .Deprecated(15, "Use FFT"); 146 REGISTER_OP("BatchIFFT") 147 .Input("input: complex64") 148 .Output("output: complex64") 149 .Deprecated(15, "Use IFFT"); 150 REGISTER_OP("BatchFFT2D") 151 .Input("input: complex64") 152 .Output("output: complex64") 153 .Deprecated(15, "Use FFT2D"); 154 REGISTER_OP("BatchIFFT2D") 155 .Input("input: complex64") 156 .Output("output: complex64") 157 .Deprecated(15, "Use IFFT2D"); 158 REGISTER_OP("BatchFFT3D") 159 .Input("input: complex64") 160 .Output("output: complex64") 161 .Deprecated(15, "Use FFT3D"); 162 REGISTER_OP("BatchIFFT3D") 163 .Input("input: complex64") 164 .Output("output: complex64") 165 .Deprecated(15, "Use IFFT3D"); 166 167 } // namespace tensorflow 168
