def __call__(self, compose_fn):
    # TODO(fengliuai): more sanity check of the input function and make sure
    # the bounded arguments of the function matches the 'inputs' and 'attrs'.
    setattr(compose_fn, '_tfr_op_name', self._op_name)

    if isinstance(func_def, types.FunctionType):
      if not hasattr(func_def, '_tfr_op_name'):
        # skip a non-composition function
        if optional:
          return (None, None)
        else:
          raise KeyError('OpDef does not exist: ' + f_name)
      op_name = getattr(func_def, '_tfr_op_name')
    elif not func_def:
      op_name = f_name
    else:

      compose_func_name.push_back('_');
      compose_func_name.push_back(tf_op_name[i] + 'a' - 'A');
    } else {
      compose_func_name.push_back(tf_op_name[i]);
    }
  }
  return compose_func_name;
}

std::string GetTFOpName(StringRef compose_func_name) {
  std::string tf_op_name;
  bool after_underscore = false;
  for (int i = 0; i < compose_func_name.size(); ++i) {

    if (!op.getCustomCode().starts_with(kFlexOpNamePrefix)) {
      return failure();
    }

    llvm::StringRef tf_op_name =
        op.getCustomCode().substr(kFlexOpNamePrefix.size());
    const std::string tf_op_full_name = llvm::Twine("tf.", tf_op_name).str();

    // Create the TF op
    OperationState op_state(op.getLoc(), tf_op_full_name);
    op_state.addOperands(op.getOperands());

// TFR function name. Examples:
//   tf.Pack => tf__pack
//   tf.ConcatV2 => tf__concat_v2
// TODO(fengliuai): move to an util file.
std::string GetComposeFuncName(StringRef tf_op_name);

// This is a hardcoded rule for mapping a TFR function op name to the
// corresponding TF opname. Examples:
//   tf__pack -> tf.Pack
//   tf__concat_v2 => tf.ConcatV2

    StringRef op_name = op->getName().stripDialect();
    if (!func_symbols_.contains(op_name)) {
      std::string tf_op_name = llvm::formatv(
          "{0}.{1}", TF::TensorFlowDialect::getDialectNamespace(), op_name);
      OperationState state =
          OperationState(loc, tf_op_name, inner_op_operands, new_types, attrs,
                         op->getSuccessors(), new_regions);
      inner_op = rewriter.create(state);

    const llvm::StringMap<Attribute>& derived_attrs) const {
  // Create the new op
  Location loc = call_op.getLoc();
  rewriter.setInsertionPointAfter(call_op);
  std::string tf_op_name = GetTFOpName(call_op.getCallee());
  OperationState new_state(loc, tf_op_name, inputs, output_types, attr_list);
  Operation* new_op = rewriter.create(new_state);
  if (materialize_derived_attrs_) {
    for (const auto& attr : derived_attrs) {