experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[tensorflow] Add a Tensor constructor that allows to transfer TensorBuffer ownership

Browse Source 
Currently Tensor does not have a constructor that allows to pass TensorBuffer ownership to Tensor. It can only be done with Ref()/Unref() operations pair.

core::RefCountPtr<TensorBuffer> clearly communicates the ownership of the buffer.

Although it would be nice to have "take ownership" constructor with a simple pointer argument, it is harder to signal this intention in a type signature.

PiperOrigin-RevId: 356993664
Change-Id: Iabe906c0c22aa4ffca0bd6d8b1fe889d286b7766
This commit is contained in:
Eugene Zhulenev 2021-02-11 09:37:27 -08:00 committed by TensorFlower Gardener
parent 2287f80894
commit 6ba169665b
2 changed files with 12 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/core/framework/tensor.cc
View File

@ -650,6 +650,12 @@ Tensor::Tensor(DataType type, const TensorShape& shape, TensorBuffer* buf)
RefIfNonNull(buf);
}
Tensor::Tensor(DataType type, const TensorShape& shape,
core::RefCountPtr<TensorBuffer> buf)
: shape_(shape), buf_(buf.release()) {
set_dtype(type);
}
bool Tensor::IsInitialized() const {
return (buf_ != nullptr && buf_->data() != nullptr) ||
shape_.num_elements() == 0;

6
tensorflow/core/framework/tensor.h
View File

@ -157,6 +157,12 @@ class Tensor {
/// Acquires a ref on buf that belongs to this Tensor.
Tensor(DataType type, const TensorShape& shape, TensorBuffer* buf);
/// \brief Creates a tensor with the input datatype, shape and buf.
///
/// Takes an ownership of the bufffer from the reference counted pointer.
Tensor(DataType type, const TensorShape& shape,
core::RefCountPtr<TensorBuffer> buf);
/// \brief Creates an empty Tensor of the given data type.
///
/// Like Tensor(), returns a 1-dimensional, 0-element Tensor with