//
// `device_name` must not name an existing physical or custom device. It must
// follow the format:
//
//    /job:<name>/replica:<replica>/task:<task>/device:<type>:<device_num>
//
// If the device is successfully registered, `status` is set to TF_OK. Otherwise
// the device is not usable. In case of a bad status, `device.delete_device` is
// still called on `device_info` (i.e. the caller does not retain ownership).

// `send_device`, `recv_device`, and `send_device_incarnation` set.
TFE_Op* SendOp(TFE_Context* ctx, TFE_TensorHandle* in,
               const std::string& op_name, const std::string& send_device,
               const std::string& recv_device,
               tensorflow::uint64 send_device_incarnation);

// Return a RecvOp op `op_name` with the attributes `send_device`,

    TF_ImportGraphDefOptions* opts, const char* prefix);

// Set the execution device for nodes in `graph_def`.
// Only applies to nodes where a device was not already explicitly specified.
// `device` is copied and has no lifetime requirements.
TF_CAPI_EXPORT extern void TF_ImportGraphDefOptionsSetDefaultDevice(
    TF_ImportGraphDefOptions* opts, const char* device);

// Set whether to uniquify imported operation names. If true, imported operation

  virtual bool UsesTFRT() = 0;

  // List attributes of available devices
  virtual void ListDevices(std::vector<DeviceAttributes>* devices) = 0;

  // Add `devices` into context's device manager. Context's device manager
  // will take ownership and maintain devices' lifetime.
  virtual Status AddDevices(std::vector<std::unique_ptr<Device>> devices) = 0;

  // Block until all pending nodes are finished.

// C API for device. The API is under active development and eventually
// should allow registering a plugin device with TensorFlow.

#ifdef __cplusplus
extern "C" {
#endif

// TF_Device is a C wrapper to the C++ TF Device class. This is to be passed
// through TF_OpKernelContext, and is opaque to plugin.
typedef struct TF_Device TF_Device;

typedef struct TF_VariableInfo TF_VariableInfo;

                                    Tensor* tensor);

absl::StatusOr<xla::PjRtCApiClient*> GetPjRtCApiClient(
    const DeviceType& device_type);

absl::Status ResetPjRtClient(const DeviceType& device_type);

}  // namespace tensorflow

// Set whether to propagate assigned device information when constructing a new
// Graph from a GraphDef. By default assigned device information is not copied
// and is re-computed by the runtime.
inline void TF_ImportGraphDefOptionsSetPropagateDeviceSpec(
    TF_ImportGraphDefOptions* opts, unsigned char propagate_device_spec) {
  opts->opts.propagate_device_spec = propagate_device_spec;
}

}  // end namespace tensorflow

  // Check if the remote task is alive.
  virtual Status CheckRemoteAlive(const std::string& remote_task_name,
                                  bool* is_alive) = 0;

  // Get pointer to the coordination service agent instance.
  virtual tsl::CoordinationServiceAgent* GetCoordinationServiceAgent() = 0;
};
}  // namespace tensorflow