From 0575b640091680cfb70f4dd93e70658de43b94f9 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Tue, 27 Jul 2021 16:45:20 -0700 Subject: [PATCH] Prevent division by 0 in LSH projection. PiperOrigin-RevId: 387225857 Change-Id: Iaeb572a763618c64f503e0026f6dd9fd769bf50c --- tensorflow/lite/kernels/lsh_projection.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorflow/lite/kernels/lsh_projection.cc b/tensorflow/lite/kernels/lsh_projection.cc index b809748c..34cbf6e5 100644 --- a/tensorflow/lite/kernels/lsh_projection.cc +++ b/tensorflow/lite/kernels/lsh_projection.cc @@ -28,7 +28,7 @@ limitations under the License. // // Input: // Tensor[0]: Hash functions. Dim.size == 2, DataType: Float. -// Tensor[0].Dim[0]: Num of hash functions. +// Tensor[0].Dim[0]: Num of hash functions. Must be at least 1. // Tensor[0].Dim[1]: Num of projected output bits generated by // each hash function. // In sparse case, Tensor[0].Dim[1] + ceil( log2(Tensor[0].Dim[0] )) <= 32. @@ -80,6 +80,7 @@ TfLiteStatus Resize(TfLiteContext* context, TfLiteNode* node) { const TfLiteTensor* input = GetInput(context, node, 1); TF_LITE_ENSURE(context, NumDimensions(input) >= 1); + TF_LITE_ENSURE(context, SizeOfDimension(input, 0) >= 1); if (NumInputs(node) == 3) { const TfLiteTensor* weight = GetInput(context, node, 2); -- 2.27.0