TF_CAPI_EXPORT extern int TF_CheckpointReaderGetVariableNumDims(
    TF_CheckpointReader* reader, const char* name);
// Load the weight of a variable
TF_CAPI_EXPORT extern TF_Tensor* TF_CheckpointReaderGetTensor(
    TF_CheckpointReader* reader, const char* name, TF_Status* status);

// TF_NewAttrBuilder() returns an object that you can set attributes on as

                                                   const char* name) {
  const auto& m = reader->GetVariableToDataTypeMap();
  return static_cast<TF_DataType>(m.at(name));
}

TF_Tensor* TF_CheckpointReaderGetTensor(TF_CheckpointReader* reader,
                                        const char* name, TF_Status* status) {
  std::unique_ptr<tensorflow::Tensor> tensor;
  reader->GetTensor(name, &tensor, status);