if (!status.ok()) {
      errors::AppendToMessage(&status, "while setting up ", DEVICE_TPU_XLA_JIT,
                              " device number ", i);
      return status;
    }
    device->SetAllowsSyncOnCompletion(false);
    if (tpu_xla_device_failure_closes_chips_flag) {
      device->SetHandleDeviceErrorCallback(&tpu::TpuNodeContext::CloseTpuHost);
    }

    devices->push_back(std::move(device));
  }

  DeviceContext* device_context;
  TF_ASSIGN_OR_RETURN(device_context, GetDeviceContextDefault());
  return MakeTensorFromProto(device_context, tensor_proto, alloc_attrs, tensor);
}

void XlaDevice::SetAllowsSyncOnCompletion(bool sync_on_completion) {
  mutex_lock lock(mu_);
  sync_on_completion_ = sync_on_completion;
}

bool XlaDevice::AllowsSyncOnCompletion() const {
  mutex_lock lock(mu_);
  return sync_on_completion_;