// Compiles a Tensorflow `function` to an HloModuleProto stored in the
  // XlaCompilationResult pointed to by `result`.
  virtual Status Compile(const XlaCompiler::CompileOptions& options,
                         const NameAttrList& function,
                         absl::Span<const XlaArgument> args,
                         XlaCompilationResult* result) = 0;

  // Compiles a Tensorflow single op to an HloModuleProto stored in the

  // Compiles a Tensorflow `function` into an HloModuleProto stored in the
  // XlaCompilationResult pointed to by `result` by calling
  // XlaCompiler::CompileFunction.
  Status Compile(const XlaCompiler::CompileOptions& options,
                 const NameAttrList& function,
                 absl::Span<const XlaArgument> args,
                 XlaCompilationResult* result) override;

    }
  })";

TEST(LegalizeMlirTest, LegalizesModule) {
  mlir::DialectRegistry mlir_registry;
  RegisterAllTensorFlowDialects(mlir_registry);

  std::vector<tensorflow::TensorShape> arg_shapes;
  XlaCompilationResult compilation_result;
  auto status = CompileSerializedMlirToXlaHlo(
      kMlirModuleStr, arg_shapes, /*device_type=*/"XLA_TPU_JIT",
      /*use_tuple_args=*/true, /*enable_op_fallback=*/false,

                                     XlaCompilationResult* result) {
  return ADD_SOURCE_LOCATION(
      xla_compiler_.CompileFunction(options, function, args, result));
}

Status TfGraphToHloCompiler::CompileSingleOp(
    const XlaCompiler::CompileOptions& options, const OpKernelContext* ctx,
    absl::Span<const XlaArgument> args, XlaCompilationResult* result) {
  return ADD_SOURCE_LOCATION(xla_compiler_.CompileSingleOp(

    } break;
    default:
      VLOG(2) << "LegalizeMlirToHlo computation input: unknown";
      break;
  }
}

Status DumpHloCompilationResult(std::string_view name,
                                XlaCompilationResult* compilation_result) {
  if (!VLOG_IS_ON(2) &&
      !DEBUG_DATA_DUMPER()->ShouldDump(std::string(name), kDebugGroupMain)) {
    return absl::OkStatus();
  }

  TF_ASSIGN_OR_RETURN(

                      llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
                          custom_legalization_passes);

// Apply shape, description, and resource information to inputs and outputs
// in the XlaCompilationResult. This should be called after
// compilation_result->computation was set.
ABSL_DEPRECATED("Not meant to be used directly and should be a util.")
Status PopulateResultIOInfo(

namespace tensorflow {
namespace tf2xla {
namespace internal {

// Legalize the given MLIR module to XLA HLO using a combination of the MLIR
// Bridge and XlaBuilder
absl::StatusOr<XlaCompilationResult> LegalizeTfToHlo(
    const tpu::MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    llvm::StringRef device_type,

    std::vector<std::vector<xla::Shape>>* per_core_arg_shapes);

// Compiles a serialized MLIR module into XLA HLO, generates all accompanying
// metadata and stores them in CompilationResult.
absl::StatusOr<XlaCompilationResult> LegalizeWithMlirBridge(
    const tpu::MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    llvm::StringRef device_type,

namespace tensorflow {
namespace tf2xla {
namespace internal {

using metrics::IncrementTfMlirBridgeSecondPhaseCounter;
using metrics::MlirBridgeSecondPhaseMetric;
using tpu::MlirToHloArgs;

absl::StatusOr<XlaCompilationResult> LegalizeTfToHlo(
    const tpu::MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    llvm::StringRef device_type,

  std::vector<std::vector<xla::Shape>> per_core_arg_shapes;
  std::vector<std::unique_ptr<mlir::Pass>> custom_legalization_passes;

  auto compilation_result = std::make_unique<XlaCompilationResult>();

  return CompileFromMlirToXlaHlo(
      compile_to_xla_hlo, mlir_to_hlo_args, metadata_proto,
      /*device_type=*/"XLA_TPU_JIT",
      /*shape_determination_fns=*/{}, use_tuple_args, compilation_result.get(),