TF_LOCKS_EXCLUDED(mu_);
  bool AllowsSyncOnCompletion() const override TF_LOCKS_EXCLUDED(mu_);

  // Installs an error handling callback when RefreshStatus sees !status.ok().
  void SetHandleDeviceErrorCallback(std::function<Status()> callback);

  Status RefreshStatus() override TF_LOCKS_EXCLUDED(mu_);

 private:
  absl::StatusOr<xla::LocalClient*> GetOrCreateClient() const;

                              " 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));
  }

  return absl::OkStatus();
}

class TpuSystemDeviceFactory : public DeviceFactory {
 public:

  sync_on_completion_ = sync_on_completion;
}

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

void XlaDevice::SetHandleDeviceErrorCallback(std::function<Status()> callback) {
  mutex_lock lock(mu_);
  device_error_callback_ = callback;
}

Status XlaDevice::HandleDeviceError() {
  std::function<Status()> local_device_error_callback;
  {