: public impl::VerifyClusteringPassBase<VerifyClusteringPass> {
 public:
  void runOnOperation() override;
};

void VerifyClusteringPass::runOnOperation() {
  Operation* func_op = getOperation();

  auto walk_result = func_op->walk([&](Operation* op) {
    if (!tensorflow::tf2xla::internal::IsInBridgeAcceptableDialects(op)) {
      std::string error = "op is in dialect " +

 public:
  void runOnOperation() override;
};

void InputMetricsLoweringPass::runOnOperation() {
  bool has_dynamic_op = false;
  Operation* func_op = getOperation();

  func_op->walk([&](Operation* op) {
    auto abstractOp = op->getRegisteredInfo();
    if (!abstractOp) return WalkResult::advance();

    if (mlir::mhlo::IsDynamicPadderOp(abstractOp->getTypeID())) {

  module_op->walk([&](func::FuncOp func_op) {
    FunctionName func_name = func_op.getSymName().str();
    existing_func_names.insert(func_name);
    // We may retrieve the original function's name from the attribute.
    // Functions without this attribute are ignored.
    auto original_func_name =
        func_op->getAttrOfType<StringAttr>("tf._original_func_name");
    if (original_func_name) {

// Returns a list of "tf.entry_function" attribute's "inputs" comma-split
// values.
//
// Example: if `func_op` has attribute `tf.entry_function = {inputs =
// "arg0:0,arg1:0"}`, then this function returns `{"arg0:0", "arg1:0"}`.
SmallVector<StringRef> GetEntryFunctionInputs(func::FuncOp func_op) {
  auto entry_function_attr =
      func_op->getAttrOfType<DictionaryAttr>("tf.entry_function");

  SmallVector<StringRef> inputs;

};

}  // namespace

void NchwConvolutionToNhwcPass::runOnOperation() {
  func::FuncOp func_op = getOperation();
  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  patterns.add<RewriteNchwConvolutionToNhwc>(&ctx);

  if (failed(applyPatternsAndFoldGreedily(func_op, std::move(patterns)))) {
    func_op.emitError() << "Failed to run NchwConvolutionToNhwcPass.";
    signalPassFailure();
  }
}

    EXPECT_TRUE(module_op_ref);
    return module_op_ref;
  }

  mlir::MLIRContext ctx_{};
};

func::FuncOp GetMainFuncOp(ModuleOp module_op) {
  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    if (func_op.getSymName() == "main") {
      return func_op;
    }
  }
  return {};
}

TEST_F(ConvertAssetArgsTest, ConvertsSingleAssetArg) {
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
    module {

};

void FoldConstantTransposePass::runOnOperation() {
  func::FuncOp func_op = getOperation();
  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  patterns.add<FoldTransposedConstantOp>(&ctx);
  if (failed(applyPatternsAndFoldGreedily(func_op, std::move(patterns)))) {
    func_op.emitError("Failed to fold constant->transpose pattern.");
    signalPassFailure();
  }
}

  return op.getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr) != nullptr;
}

void VerifyNoOutsideCompilationMarkersPass::runOnOperation() {
  Operation* func_op = getOperation();

  auto walk_result = func_op->walk([&](Operation* op) {
    if (IsDeviceClusterOp(*op) && HasChildLaunchDeviceOp(*op)) {
      std::string launch_error =
          absl::StrCat("Node `", op->getName().getStringRef().str(), "` ",

}

bool IsMhloUniformQuantizedOp(Operation& op) {
  return llvm::isa<mhlo::UniformQuantizeOp, mhlo::UniformDequantizeOp>(op);
}

void VerifyQuantLegalization::runOnOperation() {
  Operation* func_op = getOperation();

  auto walk_result = func_op->walk([&](Operation* op) {
    // Verify all uq and qint types are lowered.
    if (llvm::any_of(op->getOperandTypes(), IsQuantType) ||
        llvm::any_of(op->getResultTypes(), IsQuantType) ||

      component.Run(*module_op, quantization_config);

  EXPECT_THAT(pre_calibration_result, IsOk());

  SmallVector<func::FuncOp> func_ops;
  for (auto func_op : pre_calibration_result->getOps<func::FuncOp>()) {
    func_ops.push_back(func_op);
  }
  ASSERT_THAT(func_ops, SizeIs(2));
  EXPECT_THAT(func_ops, Contains(HasSymName("main")));
  EXPECT_THAT(func_ops, Contains(HasSymName("composite_dot_general_fn_1")));