Searched refs:host_tpu_devices (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_helpers.cc | 176 std::vector<Device*> host_tpu_devices; in GetTPUDevices() local
177 device_set.FindMatchingDevices(device_spec, &host_tpu_devices); in GetTPUDevices()
179 std::sort(host_tpu_devices.begin(), host_tpu_devices.end(), in GetTPUDevices()
186 *num_tpus_per_host = host_tpu_devices.size(); in GetTPUDevices()
187 } else if (*num_tpus_per_host != host_tpu_devices.size()) { in GetTPUDevices()
192 " vs. ", host_tpu_devices.size()); in GetTPUDevices()
194 tpu_devices->push_back(std::move(host_tpu_devices)); in GetTPUDevices()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | tpu_rewrite_device_util.cc | 132 llvm::SmallVector<Device, 8> host_tpu_devices; in GetTPUDevices() local
133 FindMatchingDevices(devices, device_spec, &host_tpu_devices); in GetTPUDevices()
136 std::sort(host_tpu_devices.begin(), host_tpu_devices.end(), in GetTPUDevices()
138 return host_tpu_devices; in GetTPUDevices()
144 auto host_tpu_devices = lookup(device); in GetTPUDevices() local
145 num_tpus_per_host = host_tpu_devices.size(); in GetTPUDevices()
146 tpu_devices->push_back(std::move(host_tpu_devices)); in GetTPUDevices()
151 auto host_tpu_devices = lookup(device_spec); in GetTPUDevices() local
153 const int64 host_tpu_devices_size = host_tpu_devices.size(); in GetTPUDevices()
157 num_tpus_per_host, ", got ", host_tpu_devices.size()); in GetTPUDevices()
[all …]
|