mlir::IRMapping &mapping) const final {
    // All tf_mlrt dialect ops can be inlined.
    return true;
  }
  // Note that CallOp and ReturnOp are handled by func; so need to implement
  // handleTerminator.
};

}  // namespace

TensorflowMlrtDialect::TensorflowMlrtDialect(mlir::MLIRContext *context)
    : mlir::Dialect(/*name=*/"tf_mlrt", context,
                    mlir::TypeID::get<TensorflowMlrtDialect>()) {

    return true;
  }

  // Handle the given inlined terminator by replacing it with a new operation
  // as necessary. Required when the region has only one block.
  void handleTerminator(Operation *op, ValueRange valuesToRepl) const final {
    auto retValOp = dyn_cast<TFRReturnOp>(op);
    if (!retValOp) return;

    for (auto ret_value : llvm::zip(valuesToRepl, retValOp.getOperands())) {