// that TF function annotated with kXlaCompiledKernelAttr (_XlaCompiledKernel).
class EncapsulateSubgraphsPass : public GraphOptimizationPass {
 public:
  Status Run(const GraphOptimizationPassOptions& options) override;
};

// A rewriting function to apply to each subgraph during encapsulation.
// 'arg_source_tensors' are the tensors corresponding to the arguments in the

  SessionOptions session_options;
  session_options.config.mutable_graph_options()
      ->mutable_optimizer_options()
      ->set_global_jit_level(OptimizerOptions::ON_2);
  GraphOptimizationPassOptions options;
  options.graph = &graph;
  options.session_options = &session_options;

  // Scope::ToGraph seems to drop assigned devices, probably because it goes

  if (global_jit_level_ == OptimizerOptions::OFF) {
    return absl::OkStatus();
  }

  return ScopingForPipelineStages();
}
}  // namespace

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

  return ClusterScopingPassImpl{graph, GetGlobalJitLevelForGraph(options)}
      .Run();
}

  for (Node* n : nodes) {
    TF_RETURN_IF_ERROR(CloneSmallConstantInputs(name_set, n));
  }
  return absl::OkStatus();
}

Status CloneConstantsForBetterClusteringPass::Run(
    const GraphOptimizationPassOptions& options) {
  if (GetGlobalJitLevelForGraph(options) == OptimizerOptions::OFF) {
    return absl::OkStatus();
  }

  Graph* g = options.graph->get();

  if (VLOG_IS_ON(1)) {

    return signalPassFailure();
  }

  // Run each of the passes that were selected.
  GraphConstructorOptions opts;
  opts.allow_internal_ops = true;
  opts.expect_device_spec = false;

  GraphOptimizationPassOptions options;
  SessionOptions sess_options;
  options.graph = &graph;
  options.flib_def = &flib_def;
  options.session_options = &sess_options;

  for (auto pass : passes_) {
    assert(pass != nullptr);

    } else {
      n->set_assigned_device_name(n->requested_device());
    }
  }

  FixupSourceAndSinkEdges(graph.get());

  GraphOptimizationPassWrapper wrapper;
  GraphOptimizationPassOptions opt_options =
      wrapper.CreateGraphOptimizationPassOptions(&graph);
  opt_options.flib_def = &flib_def;

  BuildXlaOpsPass pass(/*enable_lazy_compilation=*/true);
  TF_RETURN_IF_ERROR(pass.Run(opt_options));

  SessionOptions session_options;
  session_options.config.mutable_graph_options()
      ->mutable_optimizer_options()
      ->set_global_jit_level(OptimizerOptions::ON_2);
  GraphOptimizationPassOptions options;
  options.graph = &graph;
  options.device_set = device_set.get();
  options.session_options = &session_options;

  // Scope::ToGraph seems to drop assigned devices, probably because it goes

                           /*add_edges_to_output_of_downstream_nodes=*/true);
}

Status EncapsulateXlaComputationsPass::Run(
    const GraphOptimizationPassOptions& options) {
  VLOG(1) << "EncapsulateXlaComputations(): "
          << DumpGraphToFile("encapsulate_xla_computations_before",
                             **options.graph, options.flib_def);

    FixupSourceAndSinkEdges(g);
  }

  *changed = !slices_to_rewrite.empty();

  return absl::OkStatus();
}
}  // namespace

Status IncreaseDynamismForAutoJitPass::Run(
    const GraphOptimizationPassOptions& options) {
  MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags();
  if (flags->tf_xla_clustering_debug) {
    DumpGraphToFile("before_increase_dynamism_for_auto_jit_pass",

  for (Node* n : (*graph)->nodes()) {
    if (n->assigned_device_name().empty()) {
      n->set_assigned_device_name(kCpuDevice);
    }
  }

  GraphOptimizationPassWrapper wrapper;
  GraphOptimizationPassOptions opt_options =
      wrapper.CreateGraphOptimizationPassOptions(graph);

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

Node* FindNodeByName(const Graph& graph, const string& name) {