// A TF_AbstractTensor is an input to an operation. E.g. it could be a union
// type of eager and graph tensors. It is also the result of executing an
// operation.
typedef struct TF_AbstractTensor TF_AbstractTensor;

// A TF_AbstractOp is the metadata we need to execute an operation. E.g. this
// could contain the op type and other attributes.
typedef struct TF_AbstractOp TF_AbstractOp;

              TFE_TensorHandle* axis);

// If there is a device of type `device_type`, returns true
// and sets 'device_name' accordingly.
// `device_type` must be either "GPU" or "TPU".
bool GetDeviceName(TFE_Context* ctx, tensorflow::string* device_name,
                   const char* device_type);

// Create a ServerDef with the given `job_name` and add `num_tasks` tasks in it.

typedef jarray jintArray;
typedef jarray jlongArray;
typedef jarray jfloatArray;
typedef jarray jdoubleArray;
typedef jarray jobjectArray;

typedef jobject jweak;

// A create function. This will never actually get called in this test, it's
// just nice to know that it compiles.
void* create(TF_OpKernelConstruction* ctx) {
  TF_DataType type;
  TF_Status* s = TF_NewStatus();
  TF_OpKernelConstruction_GetAttrType(ctx, "foobar", &type, s);
  TF_DeleteStatus(s);
  return NULL;
}

// A compute function. This will never actually get called in this test, it's
// just nice to know that it compiles.

      tensorflow::PartialTensorShape* shape) const = 0;

  // Returns tensor (full) type.
  // While there is no immediate plan to deprecate dtype and shape in favor
  // of only using full type type information, this is a future possibility.
  //
  // Note that map_dtype_to_child_of_tensor() from core/framework/types.h
  // can be used to set a FullTypeDef based on dtype in a derived class if
  // appropriate.

TF_CAPI_EXPORT extern void TF_CreateAndSetPjRtCApiClient(
    const char* device_type, TF_Status* status, PJRT_NamedValue* create_options,
    int num_options);

// Resets the PjRt client for a device. After this, `TF_GetPjRtCClient` will
// returns an error for that device.
TF_CAPI_EXPORT extern void TF_ResetPjRtCClient(const char* device_type,
                                               TF_Status* status);

  // a trainable dtype.
  bool ShouldRecord(
      absl::Span<const AbstractTensorHandle* const> tensors) const;
  // Unwatches this tensor on the tape. Mainly used for cleanup when deleting
  // eager tensors.
  void DeleteTrace(const AbstractTensorHandle*);

  // Consumes the internal state of the tape (so cannot be called more than
  // once unless the tape is persistent) and produces the gradient of the target

using OwnedTFBuffer = std::unique_ptr<TF_Buffer, TFBufferDeleter>;

class CGraphOptimizer : public CustomGraphOptimizer {
 public:
  explicit CGraphOptimizer(TP_Optimizer optimizer, const char* device_type)
      : optimizer_(optimizer), device_type_(device_type) {
    if (optimizer.create_func != nullptr) {
      c_optimizer_ = (*optimizer_.create_func)();
    } else {
      c_optimizer_ = nullptr;
    }
  }

#endif  // TF_Bool

// Macro used to calculate struct size for maintaining ABI stability across
// different struct implementations.
#ifndef TF_OFFSET_OF_END
#define TF_OFFSET_OF_END(TYPE, MEMBER) \
  (offsetof(TYPE, MEMBER) + sizeof(((TYPE *)0)->MEMBER))
#endif  // TF_OFFSET_OF_END

                                    Tensor* tensor);

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

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

}  // namespace tensorflow