Class TensorReference

tensorflow/core/framework/tensor_reference.h:30–71

An opaque class that holds a reference to an underlying TensorBuffer. Unlike Tensor, it does not have any shape or type information, so it is cheaper to construct/move, but the only thing you can really do with it is Unref it, which releases one of the references to the underlying TensorBuffer. IMPORTANT: If you do not call Unref(), you will likely leak tensor memory.

// TensorBuffer.
// IMPORTANT: If you do not call Unref(), you will likely leak tensor memory.
class TensorReference {
 public:
  // Take the reference of the root buffer so the size will be more accurate
  explicit TensorReference(const Tensor& tensor);

  ~TensorReference() {}

  void Unref() const {
    if (buf_) buf_->Unref();
  }

  // Return an estimate of the total bytes being kept alive by this reference.
  size_t TotalBytes() const {
    // We add 128 as a baseline to account for per-Tensor metadata
    return 128 + (buf_ ? buf_->size() : 0);
  }

  void FillDescription(AllocationDescription* description) const {
    if (buf_) buf_->FillAllocationDescription(description);
  }

  // Convenience function for de-duplicating tensor references.
  bool SharesBufferWith(const TensorReference& t) const {
    return buf_ == t.buf_;
  }

  // Convenience function for de-duplicating tensor references.
  bool SharesBufferWith(const Tensor& t) const {
    return buf_ == (t.buf_ ? t.buf_->root_buffer() : nullptr);
  }

  // Convenience function for de-duplicating tensor references.
  size_t BufferHash() const { return std::hash<TensorBuffer*>()(buf_); }

  // A constructor used only for tests
  explicit TensorReference(TensorBuffer* test_buffer) : buf_(test_buffer) {
    if (buf_) buf_->Ref();
  }

 private:
  TensorBuffer* buf_;
};

typedef gtl::InlinedVector<TensorReference, 4> TensorReferenceVector;