Function AddOpRegistration

tensorflow/lite/interpreter_test.cc:1085–1119

Build a kernel registration for an op that copies its one input to an output

// Build a kernel registration for an op that copies its one input
// to an output
TfLiteRegistration AddOpRegistration() {
  TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};

  reg.custom_name = "my_add";
  reg.builtin_code = tflite::BuiltinOperator_CUSTOM;

  reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
    // Set output size to input size
    TfLiteTensor* input1 = &context->tensors[node->inputs->data[0]];
    TfLiteTensor* input2 = &context->tensors[node->inputs->data[1]];
    TfLiteTensor* output = &context->tensors[node->outputs->data[0]];

    TF_LITE_ENSURE_EQ(context, input1->dims->size, input2->dims->size);
    for (int i = 0; i < input1->dims->size; ++i) {
      TF_LITE_ENSURE_EQ(context, input1->dims->data[i], input2->dims->data[i]);
    }

    TF_LITE_ENSURE_STATUS(context->ResizeTensor(
        context, output, TfLiteIntArrayCopy(input1->dims)));
    return kTfLiteOk;
  };

  reg.invoke = [](TfLiteContext* context, TfLiteNode* node) {
    // Copy input data to output data.
    TfLiteTensor* a0 = &context->tensors[node->inputs->data[0]];
    TfLiteTensor* a1 = &context->tensors[node->inputs->data[1]];
    TfLiteTensor* out = &context->tensors[node->outputs->data[0]];
    int num = a0->dims->data[0];
    for (int i = 0; i < num; i++) {
      out->data.f[i] = a0->data.f[i] + a1->data.f[i];
    }
    return kTfLiteOk;
  };
  return reg;
}

class TestDelegate : public ::testing::Test {
 protected: