mlir::RegisterAllTensorFlowDialects(mlir_registry);
  mlir::MLIRContext mlir_context(mlir_registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module_with_add, mlir_module_with_sub;
  TF_ASSERT_OK(DeserializeMlirModule(module_with_add, &mlir_context,
                                     &mlir_module_with_add));
  TF_ASSERT_OK(DeserializeMlirModule(module_with_sub, &mlir_context,
                                     &mlir_module_with_sub));

  // error reporting system.
  mlir::StatusScopedDiagnosticHandler error_handler(mlir_context);

  // Parse the module.
  *mlir_module = mlir::parseSourceString<mlir::ModuleOp>(serialized_mlir_module,
                                                         mlir_context);
  if (!*mlir_module)
    return error_handler.Combine(
        errors::InvalidArgument("could not parse MLIR module"));

#include "mlir/IR/MLIRContext.h"  // from @llvm-project
#include "tensorflow/core/platform/status.h"

namespace tensorflow {

// Prints a MLIR module `module_op` and returns it as a string.
std::string SerializeMlirModule(mlir::ModuleOp module_op);

// Parses a MLIR module from `mlir_module_string` into `mlir_module` with
// context `mlir_context`.

#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace mhlo {

ConversionTarget GetDefaultLegalConversionTargets(MLIRContext& mlir_context,
                                                  bool legalize_chlo) {
  ConversionTarget target(mlir_context);

  if (legalize_chlo) {
    target.addIllegalDialect<chlo::ChloDialect>();
    target.addIllegalDialect<stablehlo::StablehloDialect>();
  } else {

TEST(DataDumperLoggerConfig, TestPassFilter) {
  mlir::DialectRegistry mlir_registry;
  mlir::RegisterAllTensorFlowDialects(mlir_registry);
  mlir::MLIRContext mlir_context(mlir_registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module_with_add;
  TF_ASSERT_OK(DeserializeMlirModule(module_with_add, &mlir_context,
                                     &mlir_module_with_add));

  std::unique_ptr<mlir::Pass> partitioning_pass =

// If legalize_chlo is true, the resulting conversion target cannot have CHLO.
mlir::ConversionTarget GetDefaultLegalConversionTargets(
    MLIRContext& mlir_context, bool legalize_chlo);

}  // namespace mhlo
}  // namespace mlir

  xla::XlaBuilder builder(
      ::testing::UnitTest::GetInstance()->current_test_info()->name());
  mlir::MLIRContext mlir_context;
  llvm::SmallVector<xla::XlaOp, 4> xla_params;
  std::vector<xla::XlaOp> returns(1);
  return BuildHloFromGraph(graph, builder, mlir_context, xla_params, returns,
                           use_output_shapes, /*args=*/{},
                           /*control_rets=*/{}, DEVICE_TPU,

namespace mlir {
namespace mhlo {
namespace test {

// Given a raw string, return a ModuleOp that can be used with the given
// MLIRContext.
absl::StatusOr<OwningOpRef<ModuleOp>> GetMlirModuleFromString(
    absl::string_view module_string, MLIRContext* mlir_context);

}  // namespace test
}  // namespace mhlo
}  // namespace mlir

  mlir::DialectRegistry mlir_registry;
  RegisterDialects(mlir_registry);
  mlir::MLIRContext mlir_context(mlir_registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module;

  TF_RETURN_IF_ERROR(
      DeserializeMlirModule(mlir_module_string, &mlir_context, &mlir_module));
  llvm::SmallVector<TensorOrResourceShape, 4> tensor_or_resource_shapes;

ABSL_DEPRECATED(
    "Use v1/compile_tf_graph.h::CompileTensorflowGraphToHlo instead.")
Status BuildHloFromGraph(
    const Graph& graph, xla::XlaBuilder& builder,
    mlir::MLIRContext& mlir_context, llvm::ArrayRef<xla::XlaOp> xla_params,
    std::vector<xla::XlaOp>& returns, bool unconditionally_use_output_shapes,
    llvm::ArrayRef<XlaArgument> args, llvm::ArrayRef<std::string> control_rets,