Method Run

tensorflow/lite/toco/graph_transformations/identify_prelu.cc:46–131

namespace toco {

::tensorflow::Status IdentifyPRelu::Run(Model* model, std::size_t op_index,
                                        bool* modified) {
  *modified = false;
  const auto add_op_it = model->operators.begin() + op_index;
  const auto* add_op = add_op_it->get();
  if (add_op == nullptr || add_op->type != OperatorType::kAdd ||
      add_op->inputs.size() != 2 ||
      add_op->fused_activation_function != FusedActivationFunctionType::kNone) {
    return ::tensorflow::Status::OK();
  }

  const auto* relu_input_op = GetOpWithOutput(*model, add_op->inputs[0]);
  if (relu_input_op == nullptr || relu_input_op->type != OperatorType::kRelu ||
      relu_input_op->inputs.size() != 1 ||
      relu_input_op->fused_activation_function !=
          FusedActivationFunctionType::kNone) {
    return ::tensorflow::Status::OK();
  }

  // TODO(ycling): Both Add and Mul are commutative. Support the case where
  // the position of operands are exchanged.
  const auto* mul_op = GetOpWithOutput(*model, add_op->inputs[1]);
  if (mul_op == nullptr || mul_op->type != OperatorType::kMul ||
      mul_op->inputs.size() != 2 ||
      mul_op->fused_activation_function != FusedActivationFunctionType::kNone) {
    return ::tensorflow::Status::OK();
  }

  const auto neg_alpha_tensor_name = mul_op->inputs[0];

  const auto* relu_neg_input_op = GetOpWithOutput(*model, mul_op->inputs[1]);

  if (relu_neg_input_op == nullptr ||
      relu_neg_input_op->inputs.size() != 1) {
    return ::tensorflow::Status::OK();
  }

  const Operator* final_input_op;
  if (relu_neg_input_op->type == OperatorType::kNeg &&
      relu_neg_input_op->fused_activation_function ==
          FusedActivationFunctionType::kRelu) {
    // This detects a Neg op with fused Relu activation function.
    final_input_op = relu_neg_input_op;
  } else {
    // This detects a Neg op followed by a separated Relu op.
    const auto* neg_input_op =
        GetOpWithOutput(*model, relu_neg_input_op->inputs[0]);
    if (neg_input_op == nullptr || neg_input_op->inputs.size() != 1 ||
        relu_neg_input_op->type != OperatorType::kRelu ||
        relu_neg_input_op->fused_activation_function !=
            FusedActivationFunctionType::kNone) {
      return ::tensorflow::Status::OK();
    }
    final_input_op = neg_input_op;
  }

  if (relu_input_op->inputs[0] != final_input_op->inputs[0]) {
    return ::tensorflow::Status::OK();