| 290 | } |
| 291 | |
| 292 | Allocator* XlaDevice::GetAllocatorLocked(AllocatorAttributes attr) { |
| 293 | if (attr.on_host()) { |
| 294 | return cpu_allocator(); |
| 295 | } |
| 296 | |
| 297 | if (xla_allocator_ == nullptr) { |
| 298 | xla::Backend* backend = client()->mutable_backend(); |
| 299 | xla_allocator_ = XlaDeviceAllocatorState::GetOrCreateXlaDeviceAllocator( |
| 300 | backend, device_ordinal_); |
| 301 | } |
| 302 | return xla_allocator_; |
| 303 | } |
| 304 | |
| 305 | Status XlaDevice::EnsureDeviceContextOk() { |
| 306 | mutex_lock lock(mu_); |
nothing calls this directly
no test coverage detected