Home
last modified time | relevance | path

Searched refs:num_devices_per_task (Results 1 – 8 of 8) sorted by relevance

/external/tensorflow/tensorflow/core/common_runtime/
Dcollective_param_resolver_local.cc453 gp->num_devices_per_task.clear(); in SetDevPerTask()
460 gp->num_devices_per_task[*last_task_name] = count; in SetDevPerTask()
465 gp->num_devices_per_task[*last_task_name] = count; in SetDevPerTask()
469 for (const auto& task_dev : gp->num_devices_per_task) { in SetDevPerTask()
Dbase_collective_executor.cc462 col_params.group.num_devices_per_task.at(task_name); in UnblockDependencies()
Dring_reducer_test.cc248 col_params_->group.num_devices_per_task[task_name] = num_devices; in Init()
/external/tensorflow/tensorflow/core/nccl/
Dcollective_communicator.cc74 const int num_local_devices = col_params->group.num_devices_per_task.at( in Enqueue()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
Dtpu_rewrite_device_util_test.cc259 int num_devices_per_task) { in MakeDeviceSet() argument
262 devices.reserve(num_tasks * num_devices_per_task + num_tasks + 1); in MakeDeviceSet()
272 for (int device = 0; device < num_devices_per_task; ++device) in MakeDeviceSet()
/external/tensorflow/tensorflow/core/framework/
Dcollective.cc69 for (const auto& dpt : num_devices_per_task) { in ToString()
Dcollective.h73 std::unordered_map<string, int32> num_devices_per_task; member
/external/tensorflow/tensorflow/core/kernels/
Dcollective_nccl_test.cc142 col_params_->group.num_devices_per_task[task_name] = num_ranks; in Init()