Fill the new `custom_initial_data(_size)?` fields in TfLiteNode.

PiperOrigin-RevId: 186389819

diff --git a/tensorflow/contrib/lite/interpreter.cc b/tensorflow/contrib/lite/interpreter.cc
index 028449211b8108d004df4d1cd8a58b4a08df6604..0c30f1c64fe6ad159acd8be85195d8e0081648df 100644 (file)
--- a/tensorflow/contrib/lite/interpreter.cc
+++ b/tensorflow/contrib/lite/interpreter.cc
@@ -25,6 +25,7 @@ limitations under the License.
 #include "tensorflow/contrib/lite/kernels/gemm_support.h"
 #include "tensorflow/contrib/lite/memory_planner.h"
 #include "tensorflow/contrib/lite/nnapi_delegate.h"
+#include "tensorflow/contrib/lite/schema/schema_generated.h"
 
 namespace {
 
@@ -298,7 +299,20 @@ TfLiteStatus Interpreter::AddNodeWithParameters(
         OpInit(*registration,
                reinterpret_cast<const char*>(builtin_data_deleter.get()), 0);
   }
+
   node.builtin_data = builtin_data_deleter.release();
+  // TODO(ycling): Filling `custom_initial_data` and `custom_initial_data_size`
+  // properly for nodes generated by ReplaceSubgraphsWithDelegateKernels.
+  if (registration->builtin_code == BuiltinOperator_CUSTOM) {
+    // When it's a CUSTOM op, the `custom_options` field in the Flatbuffer
+    // `Operator` table is passed in.
+    node.custom_initial_data = init_data;
+    node.custom_initial_data_size = init_data_size;
+  } else {
+    node.custom_initial_data = nullptr;
+    node.custom_initial_data_size = 0;
+  }
+
   node_and_reg.second = *registration;
   execution_plan_.push_back(new_node_index);
   return kTfLiteOk;