print("Allowed ModuleNames:", MODULE_CTORS.keys())
    return 1

  _, export_path, module_name = args
  module_ctor, version = MODULE_CTORS.get(module_name)
  if not module_ctor:
    print("Expected ModuleName to be one of:", MODULE_CTORS.keys())
    return 2
  os.makedirs(export_path)

  tf_module = module_ctor()
  if version == 2:

      func.return %0 : tensor<1xi32>
    }
  })";

absl::StatusOr<std::string> CompileMlirModule(bool compile_to_xla_hlo,
                                              const char* module_str) {
  MlirToHloArgs mlir_to_hlo_args;
  mlir_to_hlo_args.mlir_module = module_str;

  std::vector<TensorShape> arg_shapes;
  TPUCompileMetadataProto metadata_proto;
  bool use_tuple_args = true;
  std::vector<ShardingAndIndex> arg_core_mapping;

  if (!status.ok()) {
    return status;
  }
  return mlir_module;
}

TEST(XlaRewriteUtilTest, TestEraseClusterFuncs) {
  static const char* const module_str =
      R"(
module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:GPU:0"]} {
  func.func @convert_cluster_func(%arg0: tensor<i32>) -> () {

      enable_soft_placement,  set_original_tf_func_name};

  auto module_or = tensorflow::GraphdefToMlirTranslateFunction(
      input, input_arrays, input_dtypes, input_shapes, output_arrays,
      control_output_arrays, options, context);
  if (!module_or.status().ok()) return nullptr;
  return std::move(module_or).value();
}

static TranslateToMLIRRegistration GraphdefToMlirTranslate(

    }
  })";

absl::StatusOr<XlaCompiler::CompilationResult> CompileMlirModule(
    const char* module_str) {
  MlirToHloArgs mlir_to_hlo_args;
  mlir_to_hlo_args.rollout_state =
      ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED;
  mlir_to_hlo_args.mlir_module = module_str;

  se::Platform* platform =
      se::PlatformManager::PlatformWithName("Host").value();
  auto client =