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:

  if (import_saved_model_object_graph) {
    mlir::MLIRContext context;
    auto module_or = tensorflow::SavedModelObjectGraphToMlirImport(
        input_filename, tags, exported_names, &context);
    if (!module_or.status().ok()) return 1;

    module_or.value()->print(output->os());
  } else if (import_saved_model_signature_defs) {
    mlir::MLIRContext context;

  auto module_or = GraphdefToMlirImport(
      input, input_arrays, input_dtypes, input_shapes, output_arrays,
      control_output_arrays, import_options, context);
  if (!module_or.status().ok()) {
    LOG(ERROR) << "Graph import failed: " << module_or.status();
    return module_or.status();
  }
  auto& module = module_or.value();
  std::srand(0);

      if activation_fn is nn_ops.relu6:
        self.assertRegex(module_str, r'stablehlo.clamp.*\n.*return')
      elif activation_fn is nn_ops.relu:
        self.assertRegex(module_str, r'stablehlo.maximum.*\n.*return')
    else:
      # Check activation functions are implicit.
      self.assertNotRegex(module_str, r'stablehlo.clamp.*\n.*return')
      self.assertNotRegex(module_str, r'stablehlo.maximum.*\n.*return')

      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(

  mlir::MLIRContext context(registry);
  auto module_or = ConvertSavedModelToMlir(
      &bundle, &context, absl::Span<std::string>(exported_names));
  if (!module_or.status().ok()) {
    tsl::Set_TF_Status_from_Status(status, module_or.status());
    return "// error";
  }

  return MlirModuleToString(*std::move(module_or).value(), show_debug_info);
}

                      ->cast<mlir::MLIRContext*>();

  // Load the saved model into MLIR TF dialect.
  absl::Span<std::string> exported_names(nullptr, 0);
  auto module_or =
      tensorflow::ConvertSavedModelToMlir(&bundle, context, exported_names);
  if (!module_or.status().ok()) {
    return None();
  }

  // Make a module to wrap MLIR module and allow getting strings and running
  // transforms.

  if (saved_model_version == 2) {
    auto module_or = SavedModelObjectGraphToMlirImport(
        input_filename, tags, exported_names, context,
        /*unconditionally_use_set_output_shapes=*/true);
    if (!module_or.status().ok()) return module_or.status();
    return std::move(module_or).value();
  } else if (saved_model_version == 1) {
    MLIRImportOptions options;