namespace {

using mlir::mhlo::test::GetMlirModuleFromString;

class VerifyClusteringPassTest : public testing::Test {
 protected:
  void CreateModule(const char* module_string) {
    TF_ASSERT_OK_AND_ASSIGN(module_,
                            GetMlirModuleFromString(module_string, &context_));
    pm_ = std::make_unique<mlir::PassManager>(&context_);
    pm_->addNestedPass<mlir::func::FuncOp>(CreateVerifyClusteringPass());
  }

absl::StatusOr<OwningOpRef<ModuleOp>> GetMlirModuleFromString(
    absl::string_view module_string, MLIRContext* context) {
  mlir::DialectRegistry mlir_registry;
  RegisterAllTensorFlowDialects(mlir_registry);
  context->appendDialectRegistry(mlir_registry);

  OwningOpRef<ModuleOp> mlir_module;
  auto status =
      tensorflow::DeserializeMlirModule(module_string, context, &mlir_module);
  if (!status.ok()) {
    return status;
  }

absl::StatusOr<OwningOpRef<ModuleOp>> GetMlirModuleFromString(
    absl::string_view module_string, MLIRContext* context) {
  DialectRegistry mlir_registry;
  RegisterCommonToolingDialects(mlir_registry);
  context->appendDialectRegistry(mlir_registry);

  OwningOpRef<ModuleOp> mlir_module;
  auto status =
      tensorflow::DeserializeMlirModule(module_string, context, &mlir_module);
  if (!status.ok()) {
    return status;
  }

using ::mlir::mhlo::test::GetMlirModuleFromString;

class VerifyNoOutsideCompilationMarkersPassTest : public ::testing::Test {
 protected:
  void CreateModule(const char* module_string) {
    TF_ASSERT_OK_AND_ASSIGN(module_,
                            GetMlirModuleFromString(module_string, &context_));

    pm_ = std::make_unique<mlir::PassManager>(&context_);
    pm_->addNestedPass<func::FuncOp>(
        CreateVerifyNoOutsideCompilationMarkersPass());

    "/tensorflow/core/tf2xla/api/v2/dynamism_function_counter";

class InputLoweringMetricsPassTest : public testing::Test {
 protected:
  void CreateModule(const char* module_string) {
    TF_ASSERT_OK_AND_ASSIGN(module_,
                            GetMlirModuleFromString(module_string, &context_));
    pm_ = std::make_unique<mlir::PassManager>(&context_);
    pm_->addNestedPass<mlir::func::FuncOp>(CreateInputLoweringMetricsPass());
  }

 public:
  void SetUp() override {
    tensorflow::XlaOpRegistry::RegisterCompilationKernels();
  }

  Status CreateMlirModule(std::string module_string = kMlirModuleStr) {
    TF_ASSIGN_OR_RETURN(
        module_, test::GetMlirModuleFromString(module_string, &context_));

    context_.loadAllAvailableDialects();
    return absl::OkStatus();
  }

  Status LegalizeSingleOp(Operation& op) {

class LegalizeTfTypesTest : public Test {
 protected:
  void CreateModule(const char* module_string) {
    DialectRegistry mlir_registry;
    RegisterCommonToolingDialects(mlir_registry);
    context_.appendDialectRegistry(mlir_registry);
    TF_ASSERT_OK(
        tensorflow::DeserializeMlirModule(module_string, &context_, &module_));

    pm_ = std::make_unique<mlir::PassManager>(&context_);

class InferenceMetricsPassTest : public ::testing::Test {
 protected:
  void CreateModule(const char* module_string) {
    TF_ASSERT_OK_AND_ASSIGN(module_,
                            GetMlirModuleFromString(module_string, &context_));

    pm_ = std::make_unique<mlir::PassManager>(&context_);
    pm_->addPass(CreateInferenceMetricsPass());
  }

class VerifyTfxlaLegalizationTest : public ::testing::Test {
 protected:
  void CreateModule(const char* module_string) {
    TF_ASSERT_OK_AND_ASSIGN(module_,
                            GetMlirModuleFromString(module_string, &context_));

    pm_ = std::make_unique<mlir::PassManager>(&context_);
    pm_->addNestedPass<mlir::func::FuncOp>(

};

void MaterializePassthroughOpPass::runOnOperation() {
  getOperation().walk([](TF::MlirPassthroughOp op) {
    std::string module_string(op.getMlirModule());
    // Parse the module.
    auto nested_module =
        parseSourceString<ModuleOp>(module_string, op.getContext());
    if (!nested_module) {
      op->emitError() << "could not parse attached MLIR module";
      return;
    }
    func::FuncOp main =