}

Allocator* XlaDevice::GetAllocatorLocked(AllocatorAttributes attr) {
  if (attr.on_host()) {
    return cpu_allocator();
  }

  if (xla_allocator_ == nullptr) {
    if (UsePjRtForSingleDeviceCompilation(device_name_)) {
      VLOG(1) << "XlaDevice " << this << " uses AsyncValueAllocator";
      pjrt_allocator_ = std::make_unique<AsyncValueAllocator>();
      xla_allocator_ = pjrt_allocator_.get();
    } else {