Searched refs:gather_devices (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/python/distribute/v1/ |
| D | all_reduce.py | 558 def build_shuffle_all_reduce(input_tensors, gather_devices, red_op, un_op=None): argument
584 reduced_shards = _build_shuffle_gather(input_tensors, gather_devices,
592 def _build_shuffle_gather(input_tensors, gather_devices, red_op, un_op=None): argument
609 num_gather_devices = len(gather_devices)
620 with ops.device(gather_devices[d]):
794 def build_nccl_then_shuffle(input_tensors, gather_devices, nccl_red_op, argument
798 return build_shuffle_all_reduce(x, gather_devices, shuffle_red_op, un_op)
803 def _build_shuffle_hybrid(input_tensors, gather_devices, red_op, upper_level_f): argument
826 if len(gather_devices) != num_workers:
831 per_worker_values[w], [gather_devices[w]], red_op)
[all …]
|
| D | all_reduce_test.py | 191 gather_devices = ["/replica:0/task:0/device:CPU:0"
194 x, gather_devices, math_ops.add_n, un_op)
|