FunctionLibraryDefinition* flib_def,
                               std::unique_ptr<Graph>* result) {
  auto graph = std::make_unique<Graph>(OpRegistry::Global());
  GraphOptimizationPassOptions options;
  SessionOptions session_options;
  session_options.env = Env::Default();
  options.graph = &graph;
  options.session_options = &session_options;
  options.flib_def = flib_def;

#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"

namespace tensorflow {

Status ForceXlaConstantsOnHostPass::Run(
    const GraphOptimizationPassOptions& options) {
  Graph* graph = options.graph->get();

  OptimizerOptions opts;
  auto pflr = std::make_unique<ProcessFunctionLibraryRuntime>(
      nullptr, options.session_options->env, /*config=*/nullptr,

// -----------------------------
//
// In the future we will also translate StridedSlice and Pad a similar way.
class IncreaseDynamismForAutoJitPass : public GraphOptimizationPass {
 public:
  Status Run(const GraphOptimizationPassOptions& options) override;
};

}  // namespace tensorflow

      const DeviceSet* device_set, const ConfigProto& config_proto,
      const Graph& graph,
      const FunctionLibraryDefinition& function_library) const = 0;

  virtual Status Run(const GraphOptimizationPassOptions& options,
                     mlir::ModuleOp module) = 0;
};

class MlirV1CompatOptimizationPassRegistry {
 public:
  // Returns the global registry of MLIR optimization passes.

// which tend to be small anyway.

class CloneConstantsForBetterClusteringPass : public GraphOptimizationPass {
 public:
  CloneConstantsForBetterClusteringPass() = default;

  Status Run(const GraphOptimizationPassOptions& options) override;
};
}  // namespace tensorflow

  if (options.flib_def == nullptr) {
    return errors::InvalidArgument(
        "GraphOptimizationPassOptions::flib_def must be set for "
        "PartiallyDeclusterPass.");
  }
  if (options.session_options == nullptr ||
      options.session_options->env == nullptr) {
    return errors::InvalidArgument(
        "GraphOptimizationPassOptions::session_options::env must be set for "
        "PartiallyDeclusterPass.");
  }

  return absl::OkStatus();
}

// Replaces `n` with a `PartitionedCall` op that calls the same function.
absl::StatusOr<Node*> ReplaceFunctionCallWithPartitionedCall(
    const GraphOptimizationPassOptions& options,
    const FunctionLibraryDefinition& flib_def, Node* n, Graph* g,
    const NameAttrList& func, const Scope& root) {
  string config_string = options.session_options->config.SerializeAsString();

namespace tensorflow {

// Encapsulates nodes marked with the _xla_compile_id attribute into
// XlaLaunch operators.
class EncapsulateXlaComputationsPass : public GraphOptimizationPass {
 public:
  Status Run(const GraphOptimizationPassOptions& options) override;

  // The following methods are public only for unit tests.

  // This pass has two stages:
  // a) first, we call EncapsulateSubgraphsPass to encapsulate all nodes

  GraphOptimizationPassWrapper wrapper;
  wrapper.session_options.config.mutable_graph_options()
      ->mutable_optimizer_options()
      ->set_global_jit_level(OptimizerOptions::ON_2);
  GraphOptimizationPassOptions opt_options =
      wrapper.CreateGraphOptimizationPassOptions(graph);

  ClusterScopingPass pass;
  return pass.Run(opt_options);
}

               const Graph& graph,
               const FunctionLibraryDefinition& function_library),
              (const, override));
  MOCK_METHOD(Status, Run,
              (const GraphOptimizationPassOptions& options,
               mlir::ModuleOp module),
              (override));
};

class ModifyMlirModulePass : public MlirOptimizationPass {
 public: