ImmediateExecutionTensorHandle* input) = 0;

  virtual ImmediateExecutionContext* GetContext() const = 0;

  // Following two methods are used to support custom device.
  // Return true if the inputs contain custom device tensor handle. It means
  // that the argument need to be handled by a custom device.
  virtual bool HasCustomDeviceInput() const = 0;

  virtual const tensorflow::OpDef* OpDef() const = 0;

TF_CAPI_EXPORT extern PJRT_Client* TF_GetPjRtCClient(const char* device_type,
                                                     TF_Status* status);

// Gets the `PJRT_Buffer*` stored in the tensor. The status will contain error
// if the tensor does not have a `PjRtCApiBuffer`.
TF_CAPI_EXPORT extern PJRT_Buffer* TF_GetPjRtCBuffer(TF_Tensor* c_tensor,
                                                     TF_Status* status);

      TF_GUARDED_BY(mu);
  bool delete_requested TF_GUARDED_BY(mu);  // set true by TF_DeleteGraph

  // Used to link graphs contained in TF_WhileParams to the parent graph that
  // will eventually contain the full while loop.
  TF_Graph* parent;
  TF_Output* parent_inputs;
};

struct TF_OperationDescription {
  TF_OperationDescription(TF_Graph* g, const char* op_type,