Searched refs:reduction_op (Results 1 – 8 of 8) sorted by relevance
/external/tensorflow/tensorflow/core/nccl/ |
D | nccl_manager.cc | 136 reduction_op(reduction_op_in), in Collective() 147 const ncclRedOp_t reduction_op; // applies when <type> is a reduction. member 345 ncclRedOp_t reduction_op) { in AddToAllReduce() argument 346 AddParticipant(std::move(participant), context, kAllReduce, reduction_op); in AddToAllReduce() 370 ncclRedOp_t reduction_op) { in AddReduceSend() argument 371 AddParticipant(std::move(participant), context, kReduce, reduction_op); in AddReduceSend() 376 ncclRedOp_t reduction_op) { in AddReduceRecv() argument 377 AddParticipant(std::move(participant), context, kReduce, reduction_op); in AddReduceRecv() 398 ncclRedOp_t reduction_op) { in AddParticipant() argument 407 data_type, collective_type, reduction_op, context.num_local_devices, in AddParticipant() [all …]
|
D | nccl_manager_test.cc | 77 ncclRedOp_t reduction_op, TensorShape shape, in MakeReductionTestCase() argument 81 if (reduction_op == ncclProd) { in MakeReductionTestCase() 84 } else if (reduction_op == ncclSum) { in MakeReductionTestCase() 87 } else if (reduction_op == ncclMax) { in MakeReductionTestCase() 89 } else if (reduction_op == ncclMin) { in MakeReductionTestCase() 92 LOG(FATAL) << "Invalid reduction_op " << reduction_op; in MakeReductionTestCase() 108 if (reduction_op == ncclProd) { in MakeReductionTestCase() 110 } else if (reduction_op == ncclSum) { in MakeReductionTestCase() 112 } else if (reduction_op == ncclMax) { in MakeReductionTestCase() 116 } else if (reduction_op == ncclMin) { in MakeReductionTestCase() [all …]
|
D | nccl_manager.h | 144 const Context& context, ncclRedOp_t reduction_op); 160 const Context& context, ncclRedOp_t reduction_op); 162 const Context& context, ncclRedOp_t reduction_op); 198 ncclRedOp_t reduction_op);
|
/external/tensorflow/tensorflow/core/kernels/ |
D | collective_nccl_reducer.cc | 61 Status ReductionOp(const string& merge_op, ncclRedOp_t* reduction_op) { in ReductionOp() argument 63 *reduction_op = ncclSum; in ReductionOp() 66 *reduction_op = ncclProd; in ReductionOp() 75 ncclRedOp_t reduction_op; in Run() local 76 Status s = ReductionOp(col_params_->merge_op->type_string(), &reduction_op); in Run() 159 reduction_op); in Run()
|
D | segment_reduction_ops_test.cc | 68 std::unique_ptr<OpKernel> reduction_op( in BM_SegmentReduction() local 75 params.op_kernel = reduction_op.get(); in BM_SegmentReduction() 82 reduction_op->Compute(reduction_context.get()); in BM_SegmentReduction() 87 reduction_op->Compute(reduction_context.get()); in BM_SegmentReduction()
|
D | nccl_ops.cc | 80 ncclRedOp_t reduction_op() const { return reduction_op_; } in reduction_op() function in tensorflow::__anonc894cd950111::NcclReduceOpBase 116 reduction_op()); in ComputeAsync() 148 reduction_op()); in ComputeAsync() 185 reduction_op()); in ComputeAsync()
|
D | dynamic_partition_op_gpu.cu.cc | 400 cub::Sum reduction_op; in CountAndSortParts() local 413 num_runs_ptr, reduction_op, N, cu_stream); in CountAndSortParts() 428 num_runs_ptr, reduction_op, N, cu_stream); in CountAndSortParts()
|
/external/tensorflow/tensorflow/python/distribute/ |
D | cross_device_utils.py | 331 reduction_op='Add', argument 367 input_tensors[d], group_size, group_key, instance_key, reduction_op,
|