From 1b594a0310f9c1bc3ece2562455bba0f812efd3a Mon Sep 17 00:00:00 2001
From: Jiuqiang Tang <jqtang@google.com>
Date: Fri, 18 Nov 2022 12:42:58 -0800
Subject: [PATCH] Return error status when any tflite input and output tensor
 doesn't have valid dimensionality information that is needed to allocate
 Gl/Metal buffer before calling ModifyGraphWithDelegate.

PiperOrigin-RevId: 489539740
---
 mediapipe/calculators/tensor/BUILD                        | 2 ++
 mediapipe/calculators/tensor/inference_calculator_gl.cc   | 8 ++++++++
 .../calculators/tensor/inference_calculator_metal.cc      | 7 +++++++
 3 files changed, 17 insertions(+)

diff --git a/mediapipe/calculators/tensor/BUILD b/mediapipe/calculators/tensor/BUILD
index 4c06df0f..2a573fc4 100644
--- a/mediapipe/calculators/tensor/BUILD
+++ b/mediapipe/calculators/tensor/BUILD
@@ -464,6 +464,7 @@ cc_library(
         "//mediapipe/gpu:gl_calculator_helper",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:str_format",
         "@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
     ],
     alwayslink = 1,
@@ -513,6 +514,7 @@ cc_library(
         "//mediapipe/objc:mediapipe_framework_ios",
         "//mediapipe/util/tflite:config",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings:str_format",
         "@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate",
         "@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate_internal",
         "@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
diff --git a/mediapipe/calculators/tensor/inference_calculator_gl.cc b/mediapipe/calculators/tensor/inference_calculator_gl.cc
index bd8eb3ee..27b8bc23 100644
--- a/mediapipe/calculators/tensor/inference_calculator_gl.cc
+++ b/mediapipe/calculators/tensor/inference_calculator_gl.cc
@@ -20,6 +20,7 @@
 
 #include "absl/memory/memory.h"
 #include "absl/status/status.h"
+#include "absl/strings/str_format.h"
 #include "mediapipe/calculators/tensor/inference_calculator.h"
 #include "mediapipe/calculators/tensor/inference_calculator.pb.h"
 #include "mediapipe/framework/calculator_context.h"
@@ -154,6 +155,10 @@ absl::Status InferenceCalculatorGlImpl::GpuInferenceRunner::LoadDelegate(
   const auto& input_indices = interpreter_->inputs();
   for (int i = 0; i < input_indices.size(); ++i) {
     const TfLiteTensor* tensor = interpreter_->tensor(input_indices[i]);
+    RET_CHECK(tensor->dims->size > 0) << absl::StrFormat(
+        "Input tensor at index [%d] doesn't specify dimensions.",
+        input_indices[i]);
+
     gpu_buffers_in_.emplace_back(absl::make_unique<Tensor>(
         Tensor::ElementType::kFloat32,
         Tensor::Shape{std::vector<int>{
@@ -171,6 +176,9 @@ absl::Status InferenceCalculatorGlImpl::GpuInferenceRunner::LoadDelegate(
   // Create and bind output buffers.
   for (int i = 0; i < output_size_; ++i) {
     const TfLiteTensor* tensor = interpreter_->tensor(output_indices[i]);
+    RET_CHECK(tensor->dims->size > 0) << absl::StrFormat(
+        "Output tensor at index [%d] doesn't specify dimensions.",
+        output_indices[i]);
     gpu_buffers_out_.emplace_back(absl::make_unique<Tensor>(
         Tensor::ElementType::kFloat32,
         Tensor::Shape{std::vector<int>{
diff --git a/mediapipe/calculators/tensor/inference_calculator_metal.cc b/mediapipe/calculators/tensor/inference_calculator_metal.cc
index a85071f3..750f0456 100644
--- a/mediapipe/calculators/tensor/inference_calculator_metal.cc
+++ b/mediapipe/calculators/tensor/inference_calculator_metal.cc
@@ -22,6 +22,7 @@
 #include <vector>
 
 #include "absl/memory/memory.h"
+#include "absl/strings/str_format.h"
 #include "mediapipe/calculators/tensor/inference_calculator.h"
 #import "mediapipe/gpu/MPPMetalHelper.h"
 #include "mediapipe/gpu/MPPMetalUtil.h"
@@ -245,6 +246,9 @@ absl::Status InferenceCalculatorMetalImpl::CreateConverters(
   const auto& input_indices = interpreter_->inputs();
   for (int i = 0; i < input_indices.size(); ++i) {
     const TfLiteTensor* tensor = interpreter_->tensor(input_indices[i]);
+    RET_CHECK(tensor->dims->size > 0) << absl::StrFormat(
+        "Input tensor at index [%d] doesn't specify dimensions.",
+        input_indices[i]);
     // Create and bind input buffer.
     std::vector<int> dims{tensor->dims->data,
                           tensor->dims->data + tensor->dims->size};
@@ -266,6 +270,9 @@ absl::Status InferenceCalculatorMetalImpl::CreateConverters(
   output_shapes_.resize(output_indices.size());
   for (int i = 0; i < output_shapes_.size(); ++i) {
     const TfLiteTensor* tensor = interpreter_->tensor(output_indices[i]);
+    RET_CHECK(tensor->dims->size > 0) << absl::StrFormat(
+        "Output tensor at index [%d] doesn't specify dimensions.",
+        output_indices[i]);
     RET_CHECK(tensor->dims->size <= 4);
     // Create and bind output buffers.
     // Channels are always padded to multiple of 4.