Full set of unittests for the Eager delegate.

PiperOrigin-RevId: 206211243

2 files changed, 206 insertions, 35 deletions

diff --git a/tensorflow/contrib/lite/delegates/eager/kernel.cc b/tensorflow/contrib/lite/delegates/eager/kernel.cc
index 7235867ca5..1727981807 100644
--- a/tensorflow/contrib/lite/delegates/eager/kernel.cc
+++ b/tensorflow/contrib/lite/delegates/eager/kernel.cc
@@ -46,6 +46,28 @@ limitations under the License.
 namespace tflite {
 namespace eager {
 namespace kernel {
+
+// Controls the lifetime of tensor handles in a vector.
+class VectorOfHandles {
+ public:
+  explicit VectorOfHandles(int num_elements) : vector_(num_elements, nullptr) {}
+
+  ~VectorOfHandles() {
+    for (auto* handle : vector_) {
+      if (handle) handle->Unref();
+    }
+  }
+
+  tensorflow::gtl::InlinedVector<tensorflow::TensorHandle*, 2>* GetVector() {
+    return &vector_;
+  }
+
+  tensorflow::TensorHandle* GetHandle(int index) { return vector_[index]; }
+
+ private:
+  tensorflow::gtl::InlinedVector<tensorflow::TensorHandle*, 2> vector_;
+};
+
 // Executes the TensorFlow op given by 'op_name', with the attributes specified
 // in 'nodedef'. Inputs and outputs are given as indices into the 'buffer_map'.
 tensorflow::Status ExecuteEagerOp(tensorflow::EagerContext* eager_context,
@@ -54,8 +76,9 @@ tensorflow::Status ExecuteEagerOp(tensorflow::EagerContext* eager_context,
                                   const std::vector<int>& inputs,
                                   const std::vector<int>& outputs) {
   const tensorflow::AttrTypeMap* attr_types;
-  TF_RETURN_IF_ERROR(
-      tensorflow::AttrTypeMapForOp(op_name.c_str(), &attr_types));
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(
+      tensorflow::AttrTypeMapForOp(op_name.c_str(), &attr_types),
+      " (while processing attributes of '", op_name, "')");
 
   tensorflow::EagerOperation op(eager_context, op_name.c_str(), attr_types);
   for (const auto& attr : nodedef.attr()) {
@@ -64,7 +87,8 @@ tensorflow::Status ExecuteEagerOp(tensorflow::EagerContext* eager_context,
 
   for (int input_index : inputs) {
     if (!buffer_map->HasTensor(input_index)) {
-      return tensorflow::errors::Internal("Invalid tensor index ", input_index);
+      return tensorflow::errors::Internal(
+          "Cannot read from invalid tensor index ", input_index);
     }
     auto* handle = new tensorflow::TensorHandle(
         buffer_map->GetTensor(input_index), nullptr, nullptr, nullptr);
@@ -73,21 +97,20 @@ tensorflow::Status ExecuteEagerOp(tensorflow::EagerContext* eager_context,
   }
 
   int num_retvals = outputs.size();
-  tensorflow::gtl::InlinedVector<tensorflow::TensorHandle*, 2> retvals(
-      num_retvals, nullptr);
-
-  TF_RETURN_IF_ERROR(EagerExecute(&op, &retvals, &num_retvals));
+  VectorOfHandles retvals(num_retvals);
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(
+      EagerExecute(&op, retvals.GetVector(), &num_retvals),
+      " (while executing '", op_name, "' via Eager)");
 
-  if (outputs.size() != num_retvals) {
+  if (num_retvals != outputs.size()) {
     return tensorflow::errors::Internal(
         "Unexpected number of outputs from EagerExecute");
   }
 
   for (int i = 0; i < num_retvals; ++i) {
     const tensorflow::Tensor* tensor = nullptr;
-    TF_RETURN_IF_ERROR(retvals[i]->Tensor(&tensor));
+    TF_RETURN_IF_ERROR(retvals.GetHandle(i)->Tensor(&tensor));
     buffer_map->SetFromTensorFlow(outputs[i], *tensor);
-    retvals[i]->Unref();
   }
 
   return tensorflow::Status::OK();
@@ -145,7 +168,7 @@ void* Init(TfLiteContext* context, const char* buffer, size_t length) {
     TfLiteRegistration* reg;
     context->GetNodeAndRegistration(context, node_index, &node, &reg);
 
-    op_data->nodes.emplace_back(OpNode());
+    op_data->nodes.push_back(OpNode());
     OpNode& node_data = op_data->nodes.back();
 
     node_data.name = "";
@@ -237,7 +260,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
 
   for (auto tensor_index : op_data->subgraph_outputs) {
     if (!buffer_map->HasTensor(tensor_index)) {
-      context->ReportError(context, "Invalid tensor index %d", tensor_index);
+      context->ReportError(context, "Cannot write to invalid tensor index %d",
+                           tensor_index);
       return kTfLiteError;
     }
 
diff --git a/tensorflow/contrib/lite/delegates/eager/kernel_test.cc b/tensorflow/contrib/lite/delegates/eager/kernel_test.cc
index f0e3d432f9..7d9dddef93 100644
--- a/tensorflow/contrib/lite/delegates/eager/kernel_test.cc
+++ b/tensorflow/contrib/lite/delegates/eager/kernel_test.cc
@@ -27,9 +27,11 @@ namespace eager {
 namespace {
 
 using tensorflow::protobuf::TextFormat;
+using ::testing::ContainsRegex;
 using ::testing::ElementsAre;
 
 // We will use these are custom_names, so they need to be static.
+static const char kIdentity[] = "Identity";
 static const char kUnpack[] = "Unpack";
 static const char kAdd[] = "Add";
 static const char kMul[] = "Mul";
@@ -38,8 +40,8 @@ TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteDelegate* delegate,
                             const std::vector<int>& supported_nodes) {
   TfLiteIntArray* size_and_nodes =
       ConvertVectorToTfLiteIntArray(supported_nodes);
-  context->ReplaceSubgraphsWithDelegateKernels(context, eager::GetKernel(),
-                                               size_and_nodes, delegate);
+  TF_LITE_ENSURE_STATUS(context->ReplaceSubgraphsWithDelegateKernels(
+      context, eager::GetKernel(), size_and_nodes, delegate));
   TfLiteIntArrayFree(size_and_nodes);
   return kTfLiteOk;
 }
@@ -48,10 +50,10 @@ class KernelTest : public ::testing::Test {
  public:
   KernelTest() {
     CHECK(DelegateData::Create(&delegate_data_).ok());
-    interpreter_.reset(new Interpreter);
+    interpreter_.reset(new Interpreter(&error_reporter_));
   }
 
-  void Invoke() { ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk); }
+  bool Invoke() { return interpreter_->Invoke() == kTfLiteOk; }
 
   void SetValues(int tensor_index, const std::vector<float>& values) {
     float* v = interpreter_->typed_tensor<float>(tensor_index);
@@ -103,17 +105,18 @@ class KernelTest : public ::testing::Test {
       return " attr{ key: '" + key + "' value {" + value + "}}";
     };
 
+    string attributes;
     if (name == string(kUnpack)) {
-      string attributes = attr("T", "type: DT_FLOAT") + attr("num", "i: 2") +
-                          attr("axis", "i: 0");
-      AddTfOp(name, attributes, inputs, outputs);
+      attributes = attr("T", "type: DT_FLOAT") + attr("num", "i: 2") +
+                   attr("axis", "i: 0");
+    } else if (name == string(kIdentity)) {
+      attributes = attr("T", "type: DT_FLOAT");
     } else if (name == string(kAdd)) {
-      string attributes = attr("T", "type: DT_FLOAT");
-      AddTfOp(name, attributes, inputs, outputs);
+      attributes = attr("T", "type: DT_FLOAT");
     } else if (name == string(kMul)) {
-      string attributes = attr("T", "type: DT_FLOAT");
-      AddTfOp(name, attributes, inputs, outputs);
+      attributes = attr("T", "type: DT_FLOAT");
     }
+    AddTfOp(name, attributes, inputs, outputs);
   }
 
   void AddTensors(int num_tensors, const std::vector<int>& inputs,
@@ -121,12 +124,41 @@ class KernelTest : public ::testing::Test {
     interpreter_->AddTensors(num_tensors);
     for (int i = 0; i < num_tensors; ++i) {
       TfLiteQuantizationParams quant;
-      interpreter_->SetTensorParametersReadWrite(i, kTfLiteFloat32, /*name=*/"",
-                                                 /*dims=*/{3}, quant);
+      CHECK_EQ(interpreter_->SetTensorParametersReadWrite(i, kTfLiteFloat32,
+                                                          /*name=*/"",
+                                                          /*dims=*/{3}, quant),
+               kTfLiteOk);
     }
 
-    interpreter_->SetInputs(inputs);
-    interpreter_->SetOutputs(outputs);
+    CHECK_EQ(interpreter_->SetInputs(inputs), kTfLiteOk);
+    CHECK_EQ(interpreter_->SetOutputs(outputs), kTfLiteOk);
+  }
+
+  const TestErrorReporter& error_reporter() const { return error_reporter_; }
+
+  void AddTfLiteOp(const char* name, const std::vector<int>& inputs,
+                   const std::vector<int>& outputs) {
+    CHECK_EQ(string(name), kMul);  // can only add MUL
+    static TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
+    reg.builtin_code = BuiltinOperator_MUL;
+    reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
+      auto* i0 = &context->tensors[node->inputs->data[0]];
+      auto* o = &context->tensors[node->outputs->data[0]];
+      return context->ResizeTensor(context, o, TfLiteIntArrayCopy(i0->dims));
+    };
+    reg.invoke = [](TfLiteContext* context, TfLiteNode* node) {
+      auto* i0 = &context->tensors[node->inputs->data[0]];
+      auto* i1 = &context->tensors[node->inputs->data[1]];
+      auto* o = &context->tensors[node->outputs->data[0]];
+      for (int i = 0; i < o->bytes / sizeof(float); ++i) {
+        o->data.f[i] = i0->data.f[i] * i1->data.f[i];
+      }
+      return kTfLiteOk;
+    };
+
+    CHECK_EQ(interpreter_->AddNodeWithParameters(inputs, outputs, nullptr, 0,
+                                                 nullptr, &reg),
+             kTfLiteOk);
   }
 
  private:
@@ -151,21 +183,19 @@ class KernelTest : public ::testing::Test {
 
     flexbuffers_.push_back(fbb.GetBuffer());
     auto& buffer = flexbuffers_.back();
-    interpreter_->AddNodeWithParameters(
-        inputs, outputs, reinterpret_cast<const char*>(buffer.data()),
-        buffer.size(), nullptr, &reg);
+    CHECK_EQ(interpreter_->AddNodeWithParameters(
+                 inputs, outputs, reinterpret_cast<const char*>(buffer.data()),
+                 buffer.size(), nullptr, &reg),
+             kTfLiteOk);
   }
 
   std::unique_ptr<Interpreter> interpreter_;
   std::unique_ptr<DelegateData> delegate_data_;
   TfLiteDelegate delegate_;
   std::vector<std::vector<uint8_t>> flexbuffers_;
+  TestErrorReporter error_reporter_;
 };
 
-// TODO(ahentz): add a few more tests. In particular we need to be able to use
-// TF Lite ops along with the Eager ops, and we should check that having two or
-// more separate eager kernels (disjoint subgraphs) is OK. Also, we should be
-// verifying failure modes too.
 TEST_F(KernelTest, FullGraph) {
   // Define the graph.
   AddTensors(9, {0, 3}, {8});
@@ -187,12 +217,129 @@ TEST_F(KernelTest, FullGraph) {
   SetShape(3, {2, 2, 1});
   SetValues(3, {1.1f, 2.2f, 3.3f, 4.4f});
 
-  Invoke();
+  ASSERT_TRUE(Invoke());
 
   ASSERT_THAT(GetShape(8), ElementsAre(2, 1));
   ASSERT_THAT(GetValues(8), ElementsAre(14.52f, 38.72f));
 }
 
+TEST_F(KernelTest, BadTensorFlowOp) {
+  AddTensors(2, {0}, {1});
+  AddOp("NonExistentOp", {0}, {1});
+
+  ConfigureDelegate([](TfLiteContext* context, TfLiteDelegate* delegate) {
+    return GenericPrepare(context, delegate, {0});
+  });
+
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+
+  ASSERT_FALSE(Invoke());
+  ASSERT_THAT(error_reporter().error_messages(),
+              ContainsRegex("while processing attributes of 'NonExistentOp'"));
+}
+
+TEST_F(KernelTest, BadNumberOfOutputs) {
+  AddTensors(3, {0}, {1, 2});
+  AddOp(kIdentity, {0}, {1, 2});
+
+  ConfigureDelegate([](TfLiteContext* context, TfLiteDelegate* delegate) {
+    return GenericPrepare(context, delegate, {0});
+  });
+
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+
+  ASSERT_FALSE(Invoke());
+  ASSERT_THAT(error_reporter().error_messages(),
+              ContainsRegex("Unexpected number of outputs"));
+}
+
+TEST_F(KernelTest, IncompatibleNodeDef) {
+  AddTensors(2, {0}, {1});
+
+  // Cast is a TF op, but we don't add the proper nodedef to it in AddOp.
+  AddOp("Cast", {0}, {1});
+
+  ConfigureDelegate([](TfLiteContext* context, TfLiteDelegate* delegate) {
+    return GenericPrepare(context, delegate, {0});
+  });
+
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+
+  ASSERT_FALSE(Invoke());
+  ASSERT_THAT(error_reporter().error_messages(),
+              ContainsRegex("while executing 'Cast' via Eager"));
+}
+
+TEST_F(KernelTest, WrongSetOfNodes) {
+  AddTensors(4, {0}, {3});
+  AddOp(kUnpack, {0}, {1, 2});
+  AddTfLiteOp(kMul, {1, 2}, {3});
+
+  // Specify that kMul (#1) is supported when it actually isn't.
+  ConfigureDelegate([](TfLiteContext* context, TfLiteDelegate* delegate) {
+    return GenericPrepare(context, delegate, {0, 1});
+  });
+
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+
+  ASSERT_FALSE(Invoke());
+  ASSERT_THAT(error_reporter().error_messages(),
+              ContainsRegex("Invalid NodeDef in Eager op"));
+}
+
+TEST_F(KernelTest, MixedGraph) {
+  AddTensors(9, {0, 3}, {8});
+
+  AddOp(kUnpack, {0}, {1, 2});
+  AddOp(kUnpack, {3}, {4, 5});
+  AddOp(kAdd, {1, 4}, {6});
+  AddOp(kAdd, {2, 5}, {7});
+  AddTfLiteOp(kMul, {6, 7}, {8});
+
+  ConfigureDelegate([](TfLiteContext* context, TfLiteDelegate* delegate) {
+    return GenericPrepare(context, delegate, {0, 1, 2, 3});
+  });
+
+  SetShape(0, {2, 2, 1});
+  SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f});
+  SetShape(3, {2, 2, 1});
+  SetValues(3, {1.1f, 2.2f, 3.3f, 4.4f});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(8), ElementsAre(2, 1));
+  ASSERT_THAT(GetValues(8), ElementsAre(14.52f, 38.72f));
+}
+
+TEST_F(KernelTest, SplitGraph) {
+  AddTensors(10, {0}, {9});
+
+  AddOp(kUnpack, {0}, {1, 2});
+  AddOp(kAdd, {1, 2}, {3});
+  AddOp(kUnpack, {3}, {4, 5});
+
+  AddTfLiteOp(kMul, {4, 5}, {6});
+
+  AddOp(kUnpack, {6}, {7, 8});
+  AddOp(kAdd, {7, 8}, {9});
+
+  ConfigureDelegate([](TfLiteContext* context, TfLiteDelegate* delegate) {
+    return GenericPrepare(context, delegate, {0, 1, 2, 4, 5});
+  });
+
+  SetShape(0, {2, 2, 2, 1});
+  SetValues(0, {3.0f, 1.0f, 0.5f, -1.0f, 0.0f, 1.0f, 1.5f, 3.0f});
+
+  ASSERT_TRUE(Invoke());
+
+  ASSERT_THAT(GetShape(9), ElementsAre(1));
+  ASSERT_THAT(GetValues(9), ElementsAre(10.0f));
+}
+
 }  // namespace
 }  // namespace eager
 }  // namespace tflite