Home
last modified time | relevance | path

Searched refs:local_device_states (Results 1 – 2 of 2) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dgpu_device.cc230 std::vector<std::unique_ptr<LocalDeviceState>> local_device_states) { in BuildLocalDevices() argument
232 for (auto& local_device : local_device_states) { in BuildLocalDevices()
245 std::vector<std::unique_ptr<LocalDeviceState>> local_device_states, in BuildDistributedDevices() argument
251 for (const auto& local_device : local_device_states) { in BuildDistributedDevices()
268 std::vector<GlobalDeviceId> gpu_device_ids(local_device_states.size()); in BuildDistributedDevices()
278 local_device_states.size()); in BuildDistributedDevices()
279 TF_RET_CHECK(local_device_states[device_proto.local_device_ordinal()] != in BuildDistributedDevices()
282 std::move(local_device_states[device_proto.local_device_ordinal()]); in BuildDistributedDevices()
291 for (const auto& device : local_device_states) { in BuildDistributedDevices()
318 std::vector<std::unique_ptr<LocalDeviceState>> local_device_states, in GetGpuClient()
[all …]
Dtpu_client.cc138 std::vector<std::unique_ptr<LocalDeviceState>> local_device_states) { in GetTpuDevices() argument
163 local_device_state = std::move(local_device_states[device_ordinal]); in GetTpuDevices()
206 std::vector<std::unique_ptr<LocalDeviceState>> local_device_states; in GetTpuClient() local
207 local_device_states.reserve(client->device_count()); in GetTpuClient()
211 local_device_states.push_back( in GetTpuClient()
216 GetTpuDevices(client, std::move(local_device_states))); in GetTpuClient()