};

}  // 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) ||

  // we do not overwrite the function argument mapping here.
  func_op.walk([&](Operation* op) {
    if (op == func_op) return;
    for (Region& region : op->getRegions()) {
      assign_unique_id_to_all(region.getArguments());
    }
  });

  llvm::SmallDenseMap<ResourceHandle, int64_t> resource_handle_id_map;
  func_op.walk([&](Operation* op) {

      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")));

  explicit ConvertFuncToBfloat16Pass() = default;

 private:
  void runOnOperation() override;
};

void ConvertFuncToBfloat16Pass::runOnOperation() {
  func::FuncOp func_op = getOperation();
  MLIRContext* context = func_op.getContext();
  RewritePatternSet patterns(context);

  BFloat16TypeConverter converter;
  patterns.add<BFloat16TypePattern, BitcastConvertOpPattern>(converter,

          }
        }
      }
    });
  }
  // Remove `kSkipIslandOutlining` attributes.
  for (func::FuncOp func_op : getOperation().getOps<func::FuncOp>()) {
    if (func_op->hasAttr(mlir::TF::kSkipIslandOutlining)) {
      func_op->removeAttr(mlir::TF::kSkipIslandOutlining);
    }
  }
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>

  pm.addPass(mlir::TFDevice::CreateDecomposeResourceOpsInClusterPass());
  // Encode this in its own scope so that func_pm is not mistakenly used
  // later on.
  {
    OpPassManager& func_pm = pm.nest<FuncOp>();
    func_pm.addPass(mlir::TFTPU::CreateTPUHostComputationExpansionPass());
    func_pm.addPass(mlir::TFTPU::CreateTPUUpdateEmbeddingEnqueueOpInputsPass());
  }