Searched refs:unique_device (Results 1 – 2 of 2) sorted by relevance
123 auto unique_device = tuple_element_sharding.UniqueDevice(); in UsedDevices() local124 if (unique_device) { in UsedDevices()125 device_map[*unique_device] += 1; in UsedDevices()130 auto unique_device = UniqueDevice(); in UsedDevices() local131 if (unique_device) { in UsedDevices()132 device_map[*unique_device] += 1; in UsedDevices()249 absl::optional<int64> unique_device; in UniqueDevice() local252 if (!device || (unique_device && *device != *unique_device)) { in UniqueDevice()255 unique_device = device; in UniqueDevice()257 return unique_device; in UniqueDevice()
308 optional<int64> unique_device = batch_norm->sharding_unique_device(); in HandleBatchNormTraining() local310 unique_device.has_value() in HandleBatchNormTraining()311 ? HloSharding::AssignDevice(unique_device.value()) in HandleBatchNormTraining()412 optional<int64> unique_device = batch_norm->sharding_unique_device(); in HandleBatchNormInference() local414 unique_device.has_value() in HandleBatchNormInference()415 ? HloSharding::AssignDevice(unique_device.value()) in HandleBatchNormInference()590 auto unique_device = batch_norm->sharding_unique_device(); in HandleBatchNormGrad() local592 unique_device.has_value() in HandleBatchNormGrad()593 ? HloSharding::AssignDevice(unique_device.value()) in HandleBatchNormGrad()