tensorflow/core/tpu/graph_rewrite/tpu_embedding_software_deduplication_rewrite_pass.cc
// Inserts a SendTPUEmbeddingGradients node into the send_recv_nodes_map. This
// map temporarily holds the RecvTPUEmbeddingActivations and
// SendTPUEmbeddingGradients of the graph before they are rewritten.
Inserts a SendTPUEmbeddingGradients node into the send_recv_nodes_map. This
map temporarily holds the RecvTPUEmbeddingActivations and
SendTPUEmbeddingGradients of the graph before they are rewritten.
Inserts a SendTPUEmbeddingGradients node into the send_recv_nodes_map . This map temporarily holds the RecvTPUEmbeddingActivations and SendTPUEmbeddingGradients of the graph before they are rewritten .
Status InsertGradientsNodeIntoMap(Node* gradients_node,
                                  SendRecvNodesMap* send_recv_nodes_map) {
  std::string tpu_replicate_attr;
  TF_RETURN_IF_ERROR(GetNodeAttr(gradients_node->def(), "_tpu_replicate",
                                 &tpu_replicate_attr));
  const std::string& requested_device = gradients_node->requested_device();
  VLOG(1) << absl::StrFormat(
      "Inserting SendTPUEmbeddingGradients node with _tpu_replicate "
      "attribute: %s and requested_device: %s",
      tpu_replicate_attr, requested_device);
  const SendRecvNodesMapKey key{.tpu_replicate_attr = tpu_replicate_attr,
                                .requested_device = requested_device};
  const SendRecvNodesMap::iterator it = send_recv_nodes_map->find(key);
  if (it != send_recv_nodes_map->end()) {
    if (it->second.gradients_node != nullptr) {
      return errors::AlreadyExists(absl::StrFormat(
          "Found duplicate SendTPUEmbeddingGradients node in graph with "
          "tpu_replicate attr: %s and requested_device: %s",
          tpu_replicate_attr, requested_device));
    }
    if (it->second.activations_node == nullptr) {
      return errors::Internal(absl::StrFormat(
          "Found map object with no RecvTPUEmbeddingActivations or "
          "SendTPUEmbeddingGradients nodes and tpu_replicate attr: %s and "
          "requested_device: %s",
          tpu_replicate_attr, requested_device));
    }
    it->second.gradients_node = gradients_node;
  } else {
    send_recv_nodes_map->emplace(
        key, SendRecvNodes{.activations_node = nullptr,
                           .gradients_node = gradients_node});
  }
  return Status::OK();
}
Status InsertGradientsNodeIntoMap ( Node * gradients_node , SendRecvNodesMap * send_recv_nodes_map ) { std :: string tpu_replicate_attr ; TF_RETURN_IF_ERROR ( GetNodeAttr ( gradients_node -> def ( ) , " " , & tpu_replicate_attr ) ) ; const std :: string & requested_device = gradients_node -> requested_device ( ) ; VLOG ( 1 ) << absl :: StrFormat ( " " " " , tpu_replicate_attr , requested_device ) ; const SendRecvNodesMapKey key { . tpu_replicate_attr = tpu_replicate_attr , . requested_device = requested_device } ; const SendRecvNodesMap :: iterator it = send_recv_nodes_map -> find ( key ) ; if ( it != send_recv_nodes_map -> end ( ) ) { if ( it -> second . gradients_node != nullptr ) { return errors :: AlreadyExists ( absl :: StrFormat ( " " " " , tpu_replicate_attr , requested_device ) ) ; } if ( it -> second . activations_node == nullptr ) { return errors :: Internal ( absl :: StrFormat ( " " " " " " , tpu_replicate_attr , requested_device ) ) ; } it -> second . gradients_node = gradients_node ; } else { send_recv_nodes_map -> emplace ( key , SendRecvNodes { . activations_node = nullptr , . gradients_node = gradients_node } ) ; } return Status :: OK ( ) ; }
Inserts a SendTPUEmbeddingGradients node into the send_recv_nodes_map.
Inserts a SendTPUEmbeddingGradients node into the send_recv_nodes_map .