tensorflow/core/tpu/graph_rewrite/tpu_embedding_software_deduplication_rewrite_pass.cc
// Rewrites the graph for a particular _tpu_replicate attribute.
Rewrites the graph for a particular _tpu_replicate attribute.
Rewrites the graph for a particular _tpu_replicate attribute .
Status RewriteGraphForTpuReplicateAttrAndDevice(
    absl::string_view tpu_replicate_attr, const std::string& requested_device,
    Node* old_activations_node, Node* old_gradients_node, Graph* graph) {
  VLOG(1) << "Rewriting graph for _tpu_replicate attribute: "
          << tpu_replicate_attr << ", activations node: "
          << ((old_activations_node) ? old_activations_node->DebugString()
                                     : "NULL")
          << ", gradients node: "
          << ((old_gradients_node) ? old_gradients_node->DebugString()
                                   : "NULL");
  std::string tpu_embedding_config_str;
  TF_RETURN_IF_ERROR(ValidateAndGetTPUEmbeddingConfiguration(
      old_activations_node, old_gradients_node, tpu_replicate_attr,
      &tpu_embedding_config_str));
  Node* deduplication_data_node;
  TF_RETURN_IF_ERROR(AddRecvDeduplicationDataNode(
      old_activations_node, old_gradients_node, requested_device,
      tpu_replicate_attr, tpu_embedding_config_str, &deduplication_data_node,
      graph));
  if (old_activations_node != nullptr) {
    TF_RETURN_IF_ERROR(ReplaceRecvActivationsNodeAndAddDeduplicationInputs(
        tpu_replicate_attr, tpu_embedding_config_str, old_activations_node,
        deduplication_data_node, graph));
  }
  if (old_gradients_node != nullptr) {
    TF_RETURN_IF_ERROR(ReplaceSendGradientsNodeAndAddDeduplicationInputs(
        tpu_replicate_attr, tpu_embedding_config_str, old_gradients_node,
        deduplication_data_node, graph));
  }
  return Status::OK();
}
Status RewriteGraphForTpuReplicateAttrAndDevice ( absl :: string_view tpu_replicate_attr , const std :: string & requested_device , Node * old_activations_node , Node * old_gradients_node , Graph * graph ) { VLOG ( 1 ) << " " << tpu_replicate_attr << " " << ( ( old_activations_node ) ? old_activations_node -> DebugString ( ) : " " ) << " " << ( ( old_gradients_node ) ? old_gradients_node -> DebugString ( ) : " " ) ; std :: string tpu_embedding_config_str ; TF_RETURN_IF_ERROR ( ValidateAndGetTPUEmbeddingConfiguration ( old_activations_node , old_gradients_node , tpu_replicate_attr , & tpu_embedding_config_str ) ) ; Node * deduplication_data_node ; TF_RETURN_IF_ERROR ( AddRecvDeduplicationDataNode ( old_activations_node , old_gradients_node , requested_device , tpu_replicate_attr , tpu_embedding_config_str , & deduplication_data_node , graph ) ) ; if ( old_activations_node != nullptr ) { TF_RETURN_IF_ERROR ( ReplaceRecvActivationsNodeAndAddDeduplicationInputs ( tpu_replicate_attr , tpu_embedding_config_str , old_activations_node , deduplication_data_node , graph ) ) ; } if ( old_gradients_node != nullptr ) { TF_RETURN_IF_ERROR ( ReplaceSendGradientsNodeAndAddDeduplicationInputs ( tpu_replicate_attr , tpu_embedding_config_str , old_gradients_node , deduplication_data_node , graph ) ) ; } return Status :: OK ( ) ; }
Rewrites the graph for a particular _tpu_replicate attribute.
Rewrites the graph for a particular _tpu_replicate attribute .