From afa722d1aa5e6d846f0b036df02dd845c688bd2f Mon Sep 17 00:00:00 2001
From: "xiang.zhang" <xiang.zhang@verisilicon.com>
Date: Thu, 14 Oct 2021 19:26:05 +0800
Subject: [PATCH] Reject softmax if dim > 65535 which not supported

Signed-off-by: xiang.zhang <xiang.zhang@verisilicon.com>
---
 tensorflow/lite/delegates/nnapi/nnapi_delegate.cc | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc b/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc
index 48fa3dfd747ed6..eb6052c2de5538 100644
--- a/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc
+++ b/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc
@@ -1863,6 +1863,9 @@ bool NNAPIDelegateKernel::Validate(
       const auto& input = context->tensors[node->outputs->data[0]];
       ExpectIsFloatOrQuant8Operator(context, node, &val_ctx);
       const int input_rank = input.dims->size;
+      for(auto d : input.dims) {
+        Expect(d < 65535, NNAPIValidationFailureType::kUnsupportedOperandSize, "Input dimension larger than 65535", &val_ctx);
+      }
       Expect(input_rank <= 4,
              NNAPIValidationFailureType::kUnsupportedOperandRank,
              "Input rank should be <= 4", &val_ctx);