"tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/testdata/");
}

class TpuMetadataUtilsTest : public ::testing::Test {
 public:
  TpuMetadataUtilsTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();
  }

  absl::StatusOr<std::vector<mlir::tf_device::ClusterFuncOp>> GetClusterFuncOps(

  tensorflow::tfrt_compiler::RegisterTPULowerClusterToRuntimeOpsPassPipeline();
  tensorflow::tfrt_compiler::
      RegisterNonTPULowerClusterToRuntimeOpsPassPipeline();

  mlir::DialectRegistry registry;
  mlir::RegisterCommonToolingDialects(registry);

  return failed(
      mlir::MlirOptMain(argc, argv, "TensorFlow pass driver\n", registry));

                               std::vector<TensorShape>& arg_shapes,
                               tpu::TPUCompileMetadataProto& metadata_proto) {
  mlir::DialectRegistry registry;
  mlir::RegisterCommonToolingDialects(registry);
  mlir::MLIRContext context(registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module;

  TF_RETURN_IF_ERROR(
      DeserializeMlirModule(mlir_module_str, &context, &mlir_module));

namespace mlir::quant::stablehlo {
namespace {

std::unique_ptr<MLIRContext> CreateContext() {
  auto context = std::make_unique<MLIRContext>();
  DialectRegistry mlir_registry;
  RegisterCommonToolingDialects(mlir_registry);
  context->appendDialectRegistry(mlir_registry);
  return context;
}

TEST(IsLargeFloatTypeTest, scalars) {
  auto context = CreateContext();

    "/tensorflow/core/tf2xla/tf_quant_op_count";

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

    "/tensorflow/core/tf_mlir_bridge_first_phase_v2_count";

class FunctionClusterTensorflowDialectTest : public ::testing::Test {
 public:
  FunctionClusterTensorflowDialectTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();
  }

  absl::Status CreateMlirModule(std::string mlir_module_filename) {

    "/tensorflow/core/tf_mlir_bridge_first_phase_v2_count";

class LowerClusterToRuntimeOpsTest : public ::testing::Test {
 public:
  LowerClusterToRuntimeOpsTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();

    env_ = Env::Default();
    test_group_name_ = "TestGroup";
    test_dir_ = testing::TmpDir();

  return ::tensorflow::mangling_util::MangleTensor(tensor_proto);
}

std::unique_ptr<MLIRContext> CreateContext() {
  auto context = std::make_unique<MLIRContext>();
  DialectRegistry mlir_registry;
  RegisterCommonToolingDialects(mlir_registry);
  context->appendDialectRegistry(mlir_registry);
  context->getOrLoadDialect<tf_type::TFTypeDialect>();
  context->getOrLoadDialect<quant::QuantizationDialect>();