}

TF_Tensor* TFE_AllocateHostTensor(TFE_Context* ctx, TF_DataType dtype,
                                  const int64_t* dims, int num_dims,
                                  TF_Status* status) {
  std::vector<int64_t> dimvec(num_dims);
  for (int i = 0; i < num_dims; ++i) {
    dimvec[i] = static_cast<int64_t>(dims[i]);
  }

  if (ctx == nullptr) {
    status->status = tensorflow::errors::InvalidArgument("Invalid Context");

                   int64_t value) {
  desc->node_builder.Attr(attr_name, static_cast<int64_t>(value));
}

void TF_SetAttrIntList(TF_OperationDescription* desc, const char* attr_name,
                       const int64_t* values, int num_values) {
  desc->node_builder.Attr(
      attr_name, ArraySlice<const int64_t>(
                     reinterpret_cast<const int64_t*>(values), num_values));
}

TF_CAPI_EXPORT extern void TFE_CollectiveOpsCheckPeerHealth(
    TFE_Context* ctx, const char* task, int64_t timeout_in_ms,
    TF_Status* status);

// Information about the shape of a Tensor and its type.
struct TF_ShapeAndType {
  // Number of dimensions. -1 indicates unknown rank.
  int num_dims;
  // Array of dimensions. -1 indicates unknown dim.
  int64_t* dims;
  // The data type. May be 0 to denote unknown type.
  TF_DataType dtype;
};

                                 const int64_t* dims, int num_dims) {
  DCHECK(index >= 0 && index < shape_list->num_items);
  TF_ShapeAndType& shape = shape_list->items[index];
  DCHECK(shape.dims == nullptr) << "Shape at " << index << " is already set!";
  DCHECK(num_dims >= 0) << "Number of dimensions cannot be negative!";
  shape.num_dims = num_dims;
  shape.dims = new int64_t[num_dims];

 public:
  // Returns number of dimensions.
  virtual absl::Status NumDims(int* num_dims) const = 0;
  // Returns number of elements across all dimensions.
  virtual absl::Status NumElements(int64_t* num_elements) const = 0;
  // Returns size of specified dimension
  //
  // -1 indicates an unknown axis length; this is unreachable for most standard
  // ImmediateExecutionTensorHandles, but comes up for example when computing

namespace tensorflow {
namespace {

void TestHandleFromDLPack(TF_Status* status, TFE_Context* ctx,
                          std::vector<int64_t> shape,
                          std::vector<int64_t> strides) {
  size_t num_elements = 1;
  for (int i = 0; i < static_cast<int32_t>(shape.size()); ++i) {
    num_elements *= shape[i];
  }
  std::vector<float> data(num_elements);
  for (size_t j = 0; j < num_elements; ++j) {

TF_CAPI_EXPORT extern void TF_SetAttrInt(TF_OperationDescription* desc,
                                         const char* attr_name, int64_t value);
TF_CAPI_EXPORT extern void TF_SetAttrIntList(TF_OperationDescription* desc,
                                             const char* attr_name,
                                             const int64_t* values,
                                             int num_values);

  float A_vals[] = {1.0f, 2.0f, 3.0f, 4.0f};
  int64_t A_dims[] = {2, 2};
  AbstractTensorHandlePtr A;
  {
    AbstractTensorHandle* A_raw;
    absl::Status s = TestTensorHandleWithDims<float, TF_FLOAT>(
        ctx_.get(), A_vals, A_dims, 2, &A_raw);
    ASSERT_EQ(errors::OK, s.code()) << s.message();
    A.reset(A_raw);
  }
  float B_vals[] = {.5f, -1.0f, 1.0f, 1.0f};
  int64_t B_dims[] = {2, 2};
  AbstractTensorHandlePtr B;

  // Returns the AbstractFunction as a FunctionDef.
  virtual void GetNameAttrList(
      tensorflow::NameAttrList* name_and_attrs) const = 0;

  virtual bool GetInt(absl::string_view, int64_t* result) const = 0;
  virtual bool GetFloat(absl::string_view attr_name, float* result) const = 0;
  virtual bool GetBool(absl::string_view attr_name, bool* result) const = 0;

// gradient registerer to instantiate a GradientFunction.
struct ForwardOperation {
 public:
  string op_name;
  std::vector<AbstractTensorHandle*> inputs;
  std::vector<AbstractTensorHandle*> outputs;
  std::vector<int64_t> skip_input_indices;
  AttrBuilder attrs;
};

using GradientFunctionFactory =
    std::function<GradientFunction*(const ForwardOperation& op)>;

// Map from op name to a `GradientFunctionFactory`.