Given a "batch" dataset node, we replace the `batch_size` input with a new input that corresponds to the original input divided by `num_replicas`.
| 377 | // Given a "batch" dataset node, we replace the `batch_size` input with a new |
| 378 | // input that corresponds to the original input divided by `num_replicas`. |
| 379 | Status MutateBatchSize(const NodeDef& node, int64 num_replicas, |
| 380 | MutableGraphView* graph) { |
| 381 | // For all the batching datasets the batch_size is input number 1 except for |
| 382 | // MapAndBatchDataset. |
| 383 | int64 batch_size_arg_index = GetBatchSizeArgIndex(node); |
| 384 | NodeDef* batch_size_node = |
| 385 | graph_utils::GetInputNode(node, *graph, batch_size_arg_index); |
| 386 | int64 batch_size; |
| 387 | TF_RETURN_IF_ERROR( |
| 388 | graph_utils::GetScalarConstNodeValue(*batch_size_node, &batch_size)); |
| 389 | DCHECK_EQ(batch_size % num_replicas, 0); |
| 390 | batch_size = batch_size / num_replicas; |
| 391 | NodeDef* new_batch_size_node = |
| 392 | graph_utils::AddScalarConstNode<int64>(batch_size, graph); |
| 393 | // We don't call UpdateFanouts here because CSE elimination might lead to |
| 394 | // multiple nodes sharing the same batch size constant node. This is also |
| 395 | // why we don't delete batch_size_node as well. |
| 396 | TF_RETURN_IF_ERROR(graph->UpdateRegularFaninByPort( |
| 397 | node.name(), batch_size_arg_index, {new_batch_size_node->name(), 0})); |
| 398 | return Status::OK(); |
| 399 | } |
| 400 | |
| 401 | Status AddFlatMapNode(const string& input_dataset, |
| 402 | gtl::ArraySlice<string> other_arguments, |
no test coverage detected