file = nullptr;
    }
  }

  // The file being written to.
  std::unique_ptr<tsl::WritableFile> file;
};

// Reproducer stream that emits a reproducer to the given `llvm::raw_ostream`.
class ReproducerStream : public mlir::ReproducerStream {
 public:
  ReproducerStream(std::string name, std::unique_ptr<llvm::raw_ostream> os)
      : name_(std::move(name)), os_(std::move(os)) {}

#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/graph/graph.h"

namespace tensorflow {
namespace tf2xla {
namespace v2 {

// Given an MLIR module, returns a GraphDef.
absl::StatusOr<std::unique_ptr<GraphDef>> ConvertMlirToGraphdef(
    mlir::ModuleOp module, const GraphExportConfig& configs);

// Converts an MLIR module to TensorFlow graph and FunctionLibraryDefinition.

  tensorflow::XlaContext* context_;  // Ref-counted.

  std::unique_ptr<tensorflow::StaticDeviceMgr> device_mgr_;
  tensorflow::Device* device_;  // Owned by device_mgr_;
  std::unique_ptr<tensorflow::ScopedStepContainer> step_container_;
  std::unique_ptr<tensorflow::FunctionLibraryDefinition> flib_def_;
  std::unique_ptr<tensorflow::ProcessFunctionLibraryRuntime> pflr_;
  tensorflow::OpKernelContext::Params params_;

                                     &mlir_module_with_add));

  std::unique_ptr<mlir::Pass> partitioning_pass =
      mlir::TFTPU::CreateTPUResourceReadsWritesPartitioningPass();
  std::unique_ptr<mlir::Pass> shape_inference_pass =
      mlir::TF::CreateTFShapeInferencePass();
  std::unique_ptr<mlir::Pass> inliner_pass = mlir::createInlinerPass();

#include "tensorflow/core/lib/gtl/cleanup.h"
#include "tensorflow/core/public/session_options.h"

namespace tensorflow {
/*static*/ Status MarkForCompilationPassTestHelper::MarkForCompilation(
    std::unique_ptr<Graph>* graph, FunctionLibraryDefinition* flib_def,
    MarkForCompilationPassTestHelper::Options options) {
  // Assign all unassigned nodes to the CPU device.

// rewrite the graph to use pipelining for better device utilization.
std::unique_ptr<OperationPass<ModuleOp>> CreateEmbeddingSequencingPass();

// This is a strictly sequential and formally correct fallback option for the
// embedding pipelining pass intended for debugging during pipelining
// development.
std::unique_ptr<OperationPass<ModuleOp>> CreateEmbeddingPipeliningPass();

// Creates a pass that rewrites `tf_device.launch_func` on TPUs into TPU runtime
// ops.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateTPURewritePass(
    llvm::StringRef module_name = llvm::StringRef());

// Creates a pass that adds ops which perform formatting on variables at
// run-time according to compilation result.
std::unique_ptr<OperationPass<ModuleOp>>
CreateTPUVariableRuntimeReformattingPass();

    CHECK(DeviceFactory::AddDevices(
              SessionOptions(), "/job:localhost/replica:0/task:0", &devices_)
              .ok());
  }

 private:
  std::vector<std::unique_ptr<Device>> devices_;
};

using ::tensorflow::testing::FindNodeByName;
using ::tensorflow::testing::matchers::Attr;
using ::tensorflow::testing::matchers::CtrlDeps;
using ::tensorflow::testing::matchers::Inputs;

#include "mlir/Pass/PassInstrumentation.h"  // from @llvm-project

namespace mlir {

void RegisterPassInstrumentor(
    const std::string& name,
    std::function<std::unique_ptr<PassInstrumentation>()> creator);
std::vector<std::function<std::unique_ptr<PassInstrumentation>()>>
GetPassInstrumentors();

}  // namespace mlir

                         const SavedVariable& variable,
                         std::unique_ptr<Variable>* output);

Status LoadSavedAsset(ImmediateExecutionContext* ctx, const SavedAsset& asset,
                      const std::string& saved_model_dir,
                      absl::Span<const AssetFileDef> assets,
                      std::unique_ptr<Asset>* output);

// Creates a TFConcreteFunction from a SavedConcreteFunction.