//     subgraph3 (GPU)      subgraph4(GPU)
//         |             /
//      subgraph5 (GPU)
//         |
//      subgraph6 (CPU)
//
// The overall workflow of the pick subgraphs pass:
//  1) Build subgraphs
//    1.1) Collect output subgraphs.
//    1.2) Build `Subgraph` and their "alternative view" from FuncOp.
//  2) Pick subgraphs
//    2.1) Populate the "dp table" for (subgraph, hardware).

  // CHECK: return %[[SUBGRAPH_1]] : tensor<1024x3xf32>
  // CHECK: }
}

// -----

// Tests where StableHLO graph in main has branch.
// This test makes sure the branch will not be added to subgraph when it reaches
// a tf op:
// stablehlo.add and %0 are not in the same subgraph.

  MarkGuaranteedConstants(*graph_in_, src_arg_pairs);

  for (auto& entry : subgraphs_) {
    Subgraph& subgraph = entry.second;
    FixupSourceAndSinkEdges(subgraph.GetGraph());
  }

  if (VLOG_IS_ON(1)) {
    // Dump subgraphs.
    for (auto& entry : subgraphs_) {
      DumpGraphToFile(
          absl::StrCat("encapsulate_subgraphs_subgraph_", entry.first),

  ASSERT_THAT(status, Eq(kTfLiteOk));

  const auto subgraph = model_.subgraphs[0].get();

  // The model should only have 3 inputs and 1 output.
  EXPECT_THAT(subgraph->inputs, SizeIs(3));
  EXPECT_THAT(subgraph->outputs, SizeIs(1));

  const auto& op1 = subgraph->operators[1].get();
  const auto& op2 = subgraph->operators[2].get();
  const auto& op3 = subgraph->operators[3].get();

}

// Check if the op should be added to the subgraph.
// The op should be added to the subgraph if all of its users match one
// of following two conditions:
// 1: The user is already in the current subgraph.
// 2: The user will reach a dead end.
//
// If the op should be added to the subgraph and there are users who
// will reach the dead end, add the ops on the dead end to the subgraph as well.
bool ShouldAddOpToSubgraph(Operation* op,

    for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
         ++subgraph_idx) {
      const auto subgraph = model->subgraphs()->Get(subgraph_idx);
      for (size_t i = 0; i < subgraph->inputs()->size(); ++i) {
        if (subgraph->inputs()->Get(i) == tensor_idx) {
          return true;
        }
      }
      for (size_t i = 0; i < subgraph->outputs()->size(); ++i) {

  intermediates:[int];
}

// The root type, defining a subgraph, which typically represents an entire
// model.
table SubGraph {
  // A list of all tensors used in this subgraph.
  tensors:[Tensor];

  // Indices of the tensors that are inputs into this subgraph. Note this is
  // the list of non-static tensors that feed into the subgraph for inference.
  inputs:[int];

// to ops in the subgraph, which are not self-contained within the subgraph.
// The outputs of this function are taken to be the results of ops in the
// subgraph which are referenced as operands outside of the subgraph.
// Also refer to documention of `AccumulateOperandsDefinedAbove` &
// `AccumulateResultsDefinedWithin`.
void ExtractSubgraphToFunc(const Subgraph& subgraph, OpBuilder& builder,

        GetInferenceDeviceTypeForOp(partition_ops.front()).value();
    Subgraph old_subgraph(partition_ops, ++func_count);
    OpsAdded ops_added;
    ExtractSubgraphToFunc(old_subgraph, builder, module, ops_added);
    AddAttrs(ops_added, builder, func_count);
    // Ops in "CPU" subgraphs may nested regions with other device subgraphs.
    // We recur into these nested blocks to raise those as well. We don't raise

// RUN: tac-opt-all-backends -tfl-pick-subgraphs %s -split-input-file -verify-diagnostics | FileCheck %s

module {
  func.func @main(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>, %arg2: tensor<100xf32>, %arg3: tensor<100xf32>) -> tensor<2x100xf32> {
    %0 = func.call @func_0_GPU_FLOAT(%arg0, %arg1, %arg2) {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} : (tensor<100xf32>, tensor<100xf32>, tensor<100xf32>) -> tensor<100xf32>