Robert Griesemer <******@****.***> 1607048150 -0800

    }
    return failure();
  }
};

namespace {

bool IsShloMainFuncOp(func::FuncOp func_op) {
  if (func_op == nullptr) {
    return false;
  }

  if (!func_op.getSymName().contains(kStablehloModuleDefaultEntryFuncName)) {
    return false;
  }

  if (func_op.getSymVisibility() == "nested" ||
      func_op.getSymVisibility() == "private") {
    return false;
  }

  return true;
}

// function.
bool ShouldIncludeInMainFunction(func::FuncOp func_op) {
  return !func_op.isPrivate() && IsExported(func_op) &&
         IsEntryFunction(func_op);
}

// Sets a function to be private so it can be referred internally.
void SetFunctionPrivate(func::FuncOp func) {
  func.setVisibility(SymbolTable::Visibility::Private);

};

// Returns true iff func_op has either no Region or the body has no Blocks.
bool IsFuncOpEmpty(func::FuncOp func_op) {
  return func_op->getNumRegions() == 0 || func_op.getBody().empty();
}

// Gets the GraphOp from the function op. Returns an empty op iff it doesn't
// exist.
GraphOp GetGraphOpFromFuncOp(func::FuncOp func_op) {
  if (IsFuncOpEmpty(func_op)) return {};

      getAnalysis<TF::ResourceAliasAnalysis>();

  llvm::SmallDenseMap<func::FuncOp, PerFunctionResult, 4> per_function_results;
  llvm::SetVector<func::FuncOp> worklist;
  for (auto func_op : module.getOps<func::FuncOp>()) {
    worklist.insert(func_op);
    per_function_results.try_emplace(
        func_op, func_op, resource_alias_analysis.GetAnalysisForFunc(func_op));
  }

// exist.
// TODO(b/284222084): Move executor dialect utilities to a new library.
GraphOp GetGraphOpFromFuncOp(func::FuncOp func_op) {
  if (func_op->getNumRegions() == 0 || func_op.getBody().empty()) return {};

  auto graph_op_range = func_op.front().without_terminator();
  if (llvm::hasSingleElement(graph_op_range)) {
    // The pass runs on a valid tf_executor dialect, so the op should be the

// Returns true iff the function's symbol is public.
bool IsPublicFuncOp(func::FuncOp func_op) {
  return SymbolTable::getSymbolVisibility(&*func_op) ==
         SymbolTable::Visibility::Public;
}

}  // namespace

func::FuncOp FindMainFuncOp(ModuleOp module_op) {
  if (const auto main_func_op = module_op.lookupSymbol<func::FuncOp>(
          kImportModelDefaultGraphFuncName);

                               TF::kStablehloEntryFunctionAttrName)
                           .getValue();
  func::FuncOp func_op = symbol_table.lookup<func::FuncOp>(function_name);

  // We should not unwrap if the function is not from
  // ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass.
  if (!func_op->hasAttr(TF::kFromXlaCallModuleAttrName)) {
    return;
  }

  MLIRContext* context = call_op.getContext();

  ASSERT_TRUE(module_op);

  auto func_op = module_op->lookupSymbol<func::FuncOp>("fully_quantized_add");
  ASSERT_THAT(func_op, NotNull());

  auto add_op_itr = func_op.getBody().op_begin<mlir::stablehlo::AddOp>();
  ASSERT_THAT(add_op_itr,
              Ne(func_op.getBody().op_end<mlir::stablehlo::AddOp>()));

  EXPECT_TRUE(IsOpFullyQuantized(*add_op_itr));
}

    // Builder builder(func_op.getContext());
    func_op->setAttr(
        kTfInputShapesAttr,
        ArrayAttr::get(func_op.getContext(), updated_input_shapes_attr));
  }
}

// Validates func ops. Returns `failure` if the function is invalid.
LogicalResult ValidateFuncOp(func::FuncOp func_op) {
  auto input_shapes_attr =
      func_op->getAttrOfType<ArrayAttr>(kTfInputShapesAttr);