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

Do not write tupe index table if tuple is empty.

Browse Source 
PiperOrigin-RevId: 256438331
This commit is contained in:
A. Unique TensorFlower 2019-07-03 14:26:06 -07:00 committed by TensorFlower Gardener
parent de9c460b06
commit 43d3c42dd5
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/compiler/xla/service/transfer_manager.cc
View File

@ -244,7 +244,8 @@ Status TransferManager::WriteTupleIndexTablesAsync(
return ShapeUtil::ForEachSubshapeWithStatus( return ShapeUtil::ForEachSubshapeWithStatus(
device_buffer.on_device_shape(), device_buffer.on_device_shape(),
[&](const Shape& device_subshape, const ShapeIndex& index) -> Status { [&](const Shape& device_subshape, const ShapeIndex& index) -> Status {
if (device_subshape.IsTuple()) { if (device_subshape.IsTuple() &&
ShapeUtil::TupleElementCount(device_subshape) > 0) {
se::DeviceMemoryBase device_memory = device_buffer.buffer(index); se::DeviceMemoryBase device_memory = device_buffer.buffer(index);
TF_RET_CHECK(GetByteSizeRequirement(device_subshape) == TF_RET_CHECK(GetByteSizeRequirement(device_subshape) ==
device_memory.size()); device_memory.size());
@ -268,6 +269,9 @@ Status TransferManager::WriteTupleIndexTablesAsync(
Status TransferManager::WriteRootTupleIndexTable( Status TransferManager::WriteRootTupleIndexTable(
se::Stream* stream, const ShapedBuffer& device_buffer) { se::Stream* stream, const ShapedBuffer& device_buffer) {
TF_RET_CHECK(device_buffer.on_device_shape().IsTuple()); TF_RET_CHECK(device_buffer.on_device_shape().IsTuple());
if (ShapeUtil::TupleElementCount(device_buffer.on_device_shape()) == 0) {
return Status::OK();
}
se::DeviceMemoryBase device_memory = device_buffer.buffer({}); se::DeviceMemoryBase device_memory = device_buffer.buffer({});
TF_RET_CHECK(GetByteSizeRequirement(device_buffer.on_device_shape()) == TF_RET_CHECK(GetByteSizeRequirement(device_buffer.on_device_shape()) ==
device_memory.size()); device_memory.size());