if (e->src()->IsNextIteration()) {
          back_edges_.push_back(
              BackEdge{e, e->src(), e->src_output(), e->dst(), e->dst_input()});
        }
      }
    }
  }
  for (const BackEdge& be : back_edges_) {
    graph_->RemoveEdge(be.edge);
  }
  return absl::OkStatus();
}

const std::vector<BackEdgeHelper::BackEdge>& BackEdgeHelper::RemovedEdges()
    const {
  return back_edges_;
}

  Status Remove(Graph* graph);

  // Gets the list of removed edges.
  const std::vector<BackEdge>& RemovedEdges() const;

  // Replaces the back edges removed by a prior call to Remove.
  Status Replace();

 private:
  Graph* graph_ = nullptr;  // not owned
  std::vector<BackEdge> back_edges_;
  // Set once Replace has been called.
  bool replaced_ = false;
};

}  // namespace tensorflow

# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false -mlir-print-debuginfo %s -o - | FileCheck %s

# Verify that importing a Graph with a backedge leads to two NextIteration nodes
# to break the cycle.

# CHECK-LABEL: func @main()
# CHECK:    %[[NEXTITERATION:[a-z0-9]+]], %[[NEXTITERATION_token:[a-z0-9]+]], {{.*}} = tf_executor.NextIteration.Source
# CHECK:    tf_executor.Merge {{.*}} %[[NEXTITERATION]]

    result[i] = outputs[i].node();
  }
  return result;
}

// Manually generates the name of the `loop_var_idx`-th NextIteration node of a
// loop being constructed with `scope`. This is used to define the backedge
// before the NextIteration node is created.
string NextIterationName(const Scope& scope, int loop_var_idx) {
  string result;
  const string& prefix = scope.impl()->name();