// Fills in `dims` with the list of shapes in the attribute `attr_name` of
// `oper` and `num_dims` with the corresponding number of dimensions. On return,
// for every i where `num_dims[i]` > 0, `dims[i]` will be an array of
// `num_dims[i]` elements. A value of -1 for `num_dims[i]` indicates that the
// i-th shape in the list is unknown.
//
// The elements of `dims` will point to addresses in `storage` which must be

                           const int64_t* dims, const int num_dims,
                           const char* device, TF_Status* status) {
  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
      TFE_NewOp(context, "VarHandleOp", status), TFE_DeleteOp);
  if (TF_GetCode(status) != TF_OK) return nullptr;
  TFE_OpSetAttrType(op.get(), "dtype", type);
  TFE_OpSetAttrShape(op.get(), "shape", dims, num_dims, status);

  // Helper for constructing a resource handle and wrapping it in a `Variable`
  // object.
  static Variable* Create(TFE_Context* context, TF_DataType type,
                          const int64_t* dims, const int num_dims,
                          const char* device, TF_Status* status);
  // Dereferences the backing buffer for the variable. Note that since this can

      PartialTensorShape component_shape;
      TF_RETURN_IF_ERROR(unwrap(component.get())->Shape(&component_shape));
      if (combined_shape.dims() < 0 ||
          combined_shape.dims() != component_shape.dims()) {
        PartialTensorShape first_shape;
        TF_RETURN_IF_ERROR(unwrap(tensors_[0].get())->Shape(&first_shape));
        return errors::Unimplemented(absl::StrCat(

    auto device_value = absl::make_unique<DataType>();
    *device_value = values[device_index];
    std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)> tensor(
        TF_NewTensor(
            datatype_enum, /*dims=*/nullptr, /*num_dims=*/0,
            device_value.release(), sizeof(DataType),
            [](void* data, size_t, void* arg) {
              delete reinterpret_cast<DataType*>(data);
            },
            nullptr),

    std::unique_ptr<LoggedTensor> t, TF_Status* status) {
  auto dtype = TFE_TensorHandleDataType(t->tensor);
  TFE_CustomDeviceTensorHandleMethods handle_methods;
  handle_methods.num_dims = &LoggedTensorNumDims;
  handle_methods.dim = &LoggedTensorDim;
  handle_methods.deallocator = &LoggedTensorDeallocator;
  return TFE_NewCustomDeviceTensorHandle(context, logging_device_name.c_str(),

  // deleted, it will call ParallelTensorDeallocator to free the struct.
  ParallelTensor* t_released = t.release();
  TFE_CustomDeviceTensorHandleMethods handle_methods;
  handle_methods.num_dims = &ParallelTensorNumDims;
  handle_methods.dim = &ParallelTensorDim;
  handle_methods.deallocator = &ParallelTensorDeallocator;
  handle_methods.summarize = &ParallelTensorSummarize;