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

Internal TF change.

Browse Source 
PiperOrigin-RevId: 312780274
Change-Id: I1ec0ecccc732b04a838f2b845c0f69e4f94f5a1f
This commit is contained in:
Russell Power 2020-05-21 18:45:15 -07:00 committed by TensorFlower Gardener
parent e3290d584e
commit 86309a4f4e
2 changed files with 30 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
tensorflow/compiler/xla/literal.cc
View File

@ -2104,6 +2104,32 @@ MutableBorrowingLiteral::MutableBorrowingLiteral(const char* src_buf_ptr,
root_piece_->set_subshape(shape_.get());
}
MutableBorrowingLiteral::MutableBorrowingLiteral(absl::Span<char*> src_buf_ptrs,
const Shape& shape)
: MutableLiteralBase() {
shape_ = absl::make_unique<Shape>(shape);
if (!shape_->IsTuple()) {
CHECK_EQ(src_buf_ptrs.size(), 1);
root_piece_ = new Piece();
root_piece_->set_buffer(const_cast<char*>(src_buf_ptrs[0]));
root_piece_->set_subshape(shape_.get());
} else {
CHECK(!ShapeUtil::IsNestedTuple(*shape_));
CHECK_EQ(src_buf_ptrs.size(), ShapeUtil::TupleElementCount(*shape_));
root_piece_ = new Piece();
root_piece_->set_subshape(shape_.get());
for (int i = 0; i < src_buf_ptrs.size(); ++i) {
Piece child_piece;
const auto& src_shape = shape_->tuple_shapes(i);
CHECK(src_shape.IsArray());
child_piece.set_subshape(&src_shape);
child_piece.set_buffer(src_buf_ptrs[i]);
root_piece_->emplace_back(std::move(child_piece));
}
}
}
MutableBorrowingLiteral::~MutableBorrowingLiteral() {
if (root_piece_ != nullptr) {
delete root_piece_;

4
tensorflow/compiler/xla/literal.h
View File

@ -776,6 +776,10 @@ class MutableBorrowingLiteral : public MutableLiteralBase {
const ShapeIndex& view_root);
MutableBorrowingLiteral(const char* src_buf_ptr, const Shape& shape);
// Create a literal from a list of buffers and a shape.
// Returns a tuple literal if `shape` is a tuple type.
MutableBorrowingLiteral(absl::Span<char*> src_buf_ptrs, const Shape& shape);
private:
// Recursively copies the subtree from the `src_piece` at the given child
// index to the `dest_piece`. For buffers only the pointers are copied, but