cpovirk <******@****.***> 1721404771 -0700

  void GetTensor(const string& name,
                 std::unique_ptr<tensorflow::Tensor>* out_tensor,
                 TF_Status* out_status) const;

 private:
  // Uses "v2_reader_" to build "var name -> shape" and "var name -> data type"
  // maps; both owned by caller.
  // REQUIRES: "v2_reader_ != nullptr && v2_reader_.status().ok()".
  std::pair<std::unique_ptr<TensorSliceReader::VarToShapeMap>,

  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
  const std::vector<std::unique_ptr<ParallelTensor>>& handles = *outputs;
  const std::vector<int64_t>* shape;
  absl::Status s = handles[0]->Shape(&shape);
  ASSERT_TRUE(s.ok());
  EXPECT_EQ(0, shape->size());
}

TEST(PARALLEL_DEVICE_LIB, TestCancelOnError) {
  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
      TF_NewStatus(), TF_DeleteStatus);

namespace tensorflow {

std::string AbstractTensorHandle::DebugString() const {
  PartialTensorShape shape;
  absl::Status s = Shape(&shape);
  std::string shape_string;
  if (!s.ok()) {
    shape_string = "<error computing shape>";
  } else {
    shape_string = shape.DebugString();
  }
  return absl::StrCat("TensorHandle(shape=", shape_string,
                      ", dtype=", DataType_Name(DataType()),

namespace tensorflow {

std::string ImmediateExecutionTensorHandle::DebugString() const {
  PartialTensorShape shape;
  std::string shape_string;
  if (Shape(&shape).ok()) {
    shape_string = shape.DebugString();
  } else {
    shape_string = "<error computing shape>";
  }
  std::string value_string;
  if (!SummarizeValue(value_string).ok()) {
    value_string = "<error computing value>";
  }

  std::vector<int64_t> shape;
  int rank = -1;
  *status = handle.NumDims(&rank);
  if (!status->ok()) {
    return shape;
  }
  shape.reserve(rank);
  for (int i = 0; i < rank; ++i) {
    int64_t dim;
    *status = handle.Dim(i, &dim);
    if (!status->ok()) {
      return shape;
    }
    shape.push_back(dim);
  }
  return shape;
}

}  // namespace

extern "C" {

    ShapeHandle shape_handle = c.output(i);
    TF_ShapeAndType& shape = output_shapes_result->items[i];
    shape.num_dims = c.Rank(shape_handle);
    if (shape.num_dims == InferenceContext::kUnknownRank) {
      shape.dims = nullptr;
      continue;
    }
    shape.dims = new int64_t[shape.num_dims];
    for (size_t j = 0; j < shape.num_dims; ++j) {
      shape.dims[j] = c.Value(c.Dim(shape_handle, j));
    }
  }

  int version = TFE_CUSTOM_DEVICE_VERSION;

  // Computes the rank of the tensor handle.
  //
  // Shapes are specified via callbacks because retrieving the shape of a tensor
  // is a blocking operation for async eager; custom devices should avoid
  // retrieving shapes of tensors they wrap until the custom device tensor's
  // shape is explicitly requested where possible.
  int (*num_dims)(void* data, TF_Status* status);

  if (reader_ != nullptr) {
    status = reader_->GetTensor(name, out_tensor);
  } else {
    tensorflow::DataType dtype;
    tensorflow::TensorShape shape;
    status = v2_reader_->LookupDtypeAndShape(name, &dtype, &shape);
    if (status.ok()) {
      out_tensor->reset(new Tensor(dtype, shape));
      status = v2_reader_->Lookup(name, out_tensor->get());
      if (!status.ok()) out_tensor->reset();
    }
  }
  if (!status.ok()) {

  dltensor_in->device = {kDLCPU, 0};
  dltensor_in->ndim = static_cast<int32_t>(shape.size());
  dltensor_in->dtype = {kDLFloat, 32, 1};
  dltensor_in->shape = shape.data();
  dltensor_in->strides = strides.data();
  TFE_TensorHandle* handle = TFE_HandleFromDLPack(&dlm_in, status, ctx);
  ASSERT_NE(handle, nullptr)
      << TF_Message(status) << " (shape=[" << absl::StrJoin(shape, ",")
      << "], strides=[" << absl::StrJoin(strides, ",") << "])";