tensorflow/compiler/tf2xla/rearrange_function_argument.cc
// For StatefulPartitionedCall node, given mapping between original input index
// and rearranged input index, reorder output edges for the node. DT_RESOURCE
// outputs are removed from the node and we will use the node's corresponding
// input for the edge.
For StatefulPartitionedCall node, given mapping between original input index
and rearranged input index, reorder output edges for the node. DT_RESOURCE
outputs are removed from the node and we will use the node's corresponding
input for the edge.
For StatefulPartitionedCall node given mapping between original input index and rearranged input index reorder output edges for the node . DT_RESOURCE outputs are removed from the node and we will use the node ' s corresponding input for the edge .
Status RearrangeOutputEdges(Node* n, Graph* g,
                            const std::map<int, int>& retval_index_mapping,
                            const std::map<int, int>& resource_retval_to_arg) {
  std::vector<const Edge*> out_edges;
  for (const Edge* e : n->out_edges()) {
    if (!e->IsControlEdge()) {
      out_edges.push_back(e);
    }
  }
  for (const Edge* e : out_edges) {
    Node* dst = e->dst();
    int dst_input = e->dst_input();
    int src_output = e->src_output();
    auto iter = retval_index_mapping.find(src_output);
    if (iter == retval_index_mapping.end()) {
      TF_RET_CHECK(resource_retval_to_arg.find(src_output) !=
                   resource_retval_to_arg.end());
      g->RemoveEdge(e);
      const Edge* input_edge;
      TF_RETURN_IF_ERROR(
          n->input_edge(resource_retval_to_arg.at(src_output), &input_edge));
      g->AddEdge(input_edge->src(), input_edge->src_output(), dst, dst_input);
    } else {
      g->RemoveEdge(e);
      g->AddEdge(n, iter->second, dst, dst_input);
    }
  }
  return Status::OK();
}
Status RearrangeOutputEdges ( Node * n , Graph * g , const std :: map < int , int > & retval_index_mapping , const std :: map < int , int > & resource_retval_to_arg ) { std :: vector < const Edge * > out_edges ; for ( const Edge * e : n -> out_edges ( ) ) { if ( ! e -> IsControlEdge ( ) ) { out_edges . push_back ( e ) ; } } for ( const Edge * e : out_edges ) { Node * dst = e -> dst ( ) ; int dst_input = e -> dst_input ( ) ; int src_output = e -> src_output ( ) ; auto iter = retval_index_mapping . find ( src_output ) ; if ( iter == retval_index_mapping . end ( ) ) { TF_RET_CHECK ( resource_retval_to_arg . find ( src_output ) != resource_retval_to_arg . end ( ) ) ; g -> RemoveEdge ( e ) ; const Edge * input_edge ; TF_RETURN_IF_ERROR ( n -> input_edge ( resource_retval_to_arg . at ( src_output ) , & input_edge ) ) ; g -> AddEdge ( input_edge -> src ( ) , input_edge -> src_output ( ) , dst , dst_input ) ; } else { g -> RemoveEdge ( e ) ; g -> AddEdge ( n , iter -> second , dst , dst_input ) ; } } return Status :: OK ( ) ; }
For StatefulPartitionedCall node, given mapping between original input index
and rearranged input index, reorder output edges for the node.
For StatefulPartitionedCall node given mapping between original input index and rearranged input index reorder output edges for the node .