Searched refs:unshard_op_node (Results 1 – 1 of 1) sorted by relevance
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | host_training_loop_optimization_util.cc | 592 Node* unshard_op_node = graph->AddNode(unshard_node_def, &status); in AddReshardOp() local
595 unshard_op_node->set_assigned_device_name( in AddReshardOp()
608 graph->AddEdge(enter_node, 0, unshard_op_node, i); in AddReshardOp()
612 graph->AddControlEdge(after_last_iteration_node, unshard_op_node); in AddReshardOp()
613 graph->AddControlEdge(unshard_op_node, after_unshard_node); in AddReshardOp()
615 graph->AddEdge(default_sharding_node, 0, unshard_op_node, new_key_input); in AddReshardOp()
618 graph->AddEdge(var_handle_node, 0, unshard_op_node, format_state_input); in AddReshardOp()
|