Searched refs:system_devices (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | tpu_rewrite_device_util.cc | 91 llvm::SmallVector<Device, 8> system_devices; in GetTPUSystemDevices() local
92 FindMatchingDevices(devices, spec, &system_devices); in GetTPUSystemDevices()
93 if (system_devices.empty()) in GetTPUSystemDevices()
97 const auto& job = system_devices[0].job; in GetTPUSystemDevices()
98 auto replica = system_devices[0].replica; in GetTPUSystemDevices()
99 for (const auto& device : llvm::make_range(std::next(system_devices.begin()), in GetTPUSystemDevices()
100 system_devices.end())) { in GetTPUSystemDevices()
110 std::sort(system_devices.begin(), system_devices.end(), in GetTPUSystemDevices()
113 matched_devices->swap(system_devices); in GetTPUSystemDevices()
122 Devices devices, llvm::ArrayRef<Device> system_devices, in GetTPUDevices() argument
[all …]
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_helpers.cc | 46 std::vector<Device*> system_devices; in GetSystemDevice() local
47 device_set.FindMatchingDevices(*system_spec, &system_devices); in GetSystemDevice()
48 if (system_devices.empty()) { in GetSystemDevice()
56 } else if (system_devices.size() > 1) { in GetSystemDevice()
59 for (auto device : system_devices) { in GetSystemDevice()
74 std::sort(system_devices.begin(), system_devices.end(), in GetSystemDevice()
85 *system_device = system_devices[0]; in GetSystemDevice()
|