// Similar to TF_GraphToFunction but allows specifying control outputs of the
// function.
//
//  The arguments of TF_GraphToFunction have the same meaning, but the new
//  arguments are as follows:
//
//    ncontrol_outputs: Number of control outputs of the function.
//    control_outputs: vector of TF_Operation objects to be marked as control

// devices, so executing tf.functions which contain operations placed on the
// custom devices will fail.
//
// `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

// `enable_xla_compilation` is non-zero, and OFF otherwise.
// b) ConfigProto.gpu_options.allow_growth is set to `gpu_memory_allow_growth`.
// c) ConfigProto.device_count is set to `num_cpu_devices`.
TF_CAPI_EXPORT extern TF_Buffer* TF_CreateConfig(
    unsigned char enable_xla_compilation, unsigned char gpu_memory_allow_growth,
    unsigned int num_cpu_devices);

//  - target: dy in dx/dy
//  - source: dx in dx/dy
//  - tensor: one of the many inputs or outputs of an operation
//
// Below here we do the gradient algorithm. It works as follows:
//
// First we filter the tape to just the subset of operations we want to
// differentiate. In the process of doing so we count how many times each Tensor

#include "tensorflow/c/c_api_macros.h"
#include "tensorflow/c/tf_status.h"

// --------------------------------------------------------------------------
// C API for Graph. The API is under active development and eventually
// should allow registering a plugin graph optimizer with TensorFlow.
//
// Conventions:
//   * Struct prefix indicates whether struct fields should be filled by the
//     plugin or core implementation:

typedef void (*TFInitGraphPluginFn)(TP_OptimizerRegistrationParams* const,
                                    TF_Status* const);

// Registers Graph optimizers.
Status InitGraphPlugin(void* dso_handle);

// Allow registering a graph optimizer using a function (used for
// testing).
Status InitGraphPlugin(TFInitGraphPluginFn init_fn);

struct GrapplerItem;
class Cluster;

struct TFStatusDeleter {

/// the 3 types of files. In all other cases, there is no need to modify the
/// address of the opaque data pointer, hence the wrapper pointer is marked
/// `const`.
///
/// For consistency, the arguments on all these functions follow the same
/// pattern: first we have the opaque pointer argument ("this" above), then the
/// input arguments, then the in-out arguments (if any) and we finish the

#include "tensorflow/c/tf_status.h"
#include "xla/pjrt/c/pjrt_c_api.h"

// --------------------------------------------------------------------------
// 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.

// maintains the map from `op_id` to a `OpTapeEntry` which stores the `op_type`,
// inputs and outputs and the gradient function These data structures combined
// allow us to trace the data dependencies between operations and hence compute
// gradients.
//
// `ZerosLike` is not expected to be called and returns a nullptr. The creation