From 1c56f53be0b722ca657cbc7df461ed676c8642a2 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Tue, 4 May 2021 17:11:46 -0700 Subject: [PATCH] Fix a check fail in Fast Fourier implementation PiperOrigin-RevId: 372026629 Change-Id: Id05c3362aa575271bc3e06b16316c9037085fc11 --- tensorflow/core/kernels/fft_ops.cc | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tensorflow/core/kernels/fft_ops.cc b/tensorflow/core/kernels/fft_ops.cc index c3235c1547ff0..29ca0d2f546a0 100644 --- a/tensorflow/core/kernels/fft_ops.cc +++ b/tensorflow/core/kernels/fft_ops.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "tensorflow/core/platform/errors.h" #define EIGEN_USE_THREADS // See docs in ../ops/fft_ops.cc. @@ -261,6 +262,9 @@ class FFTCPU : public FFTBase { i == FFTRank ? fft_shape[i - 1] / 2 + 1 : fft_shape[i - 1]; full_fft_shape.AddDim(fft_shape[i - 1]); } + OP_REQUIRES(ctx, full_fft_shape.num_elements() > 0, + errors::InvalidArgument("Obtained a FFT shape of 0 elements: ", + full_fft_shape.DebugString())); Tensor temp; OP_REQUIRES_OK(ctx, ctx->allocate_temp(DataTypeToEnum::v(),