expected_max_outputs, cancellation_manager, step_id);
  }
}

void ParallelDevice::AsyncWait(TFE_Context* context, TF_Status* status) const {
  StatusPtr first_bad_status(nullptr);

  for (const auto& dt : device_threads_) {
    StatusPtr async_wait_status(TF_NewStatus());
    dt->AsyncWait(async_wait_status.get());
    // Prefer non cancelled errors to uncover real failures.

  virtual absl::Status AddDevices(
      std::vector<std::unique_ptr<Device>> devices) = 0;

  // Block until all pending nodes are finished.
  virtual absl::Status AsyncWait() = 0;

  // Add a function (serialized FunctionDef protocol buffer) so that it can
  // be executed as an op. Return error if the function with the same name
  // already exists.

  absl::optional<std::vector<std::unique_ptr<ParallelTensor>>> Join(
      const std::vector<PartialTensorShape>& expected_output_shapes,
      TF_Status* status) const;

  void AsyncWait(TFE_Context* context, TF_Status* status) const;

  // Device strings for component devices that only include a
  // worker/task/replica if any of those differ across components. Useful for
  // printing debug messages.

#if defined(IS_MOBILE_PLATFORM)
  status->status = tensorflow::OkStatus();
#else   // !defined(IS_MOBILE_PLATFORM)
  status->status = tensorflow::unwrap(ctx)->AsyncWait();
#endif  // !IS_MOBILE_PLATFORM
}

void TFE_ContextSetThreadLocalDevicePlacementPolicy(
    TFE_Context* ctx, TFE_ContextDevicePlacementPolicy policy) {