Searched refs:device_subshape (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | generic_transfer_manager.cc | 110 [&](const Shape& device_subshape, const ShapeIndex& index) -> Status { in TransferLiteralToDeviceAsync() argument
112 if (device_subshape.IsArray()) { in TransferLiteralToDeviceAsync()
113 TF_RET_CHECK(GetByteSizeRequirement(device_subshape) == in TransferLiteralToDeviceAsync()
119 if (LayoutUtil::Equal(device_subshape.layout(), in TransferLiteralToDeviceAsync()
124 /*size=*/GetByteSizeRequirement(device_subshape), source, in TransferLiteralToDeviceAsync()
128 relayed_out_literal = subliteral.Relayout(device_subshape.layout(), in TransferLiteralToDeviceAsync()
133 /*size=*/GetByteSizeRequirement(device_subshape), source, in TransferLiteralToDeviceAsync()
|
| D | transfer_manager.cc | 299 [&](const Shape& device_subshape, const ShapeIndex& index) -> Status { in WriteTupleIndexTablesAsync() argument
300 if (device_subshape.IsTuple() && in WriteTupleIndexTablesAsync()
301 ShapeUtil::TupleElementCount(device_subshape) > 0) { in WriteTupleIndexTablesAsync()
303 TF_RET_CHECK(GetByteSizeRequirement(device_subshape) == in WriteTupleIndexTablesAsync()
308 for (int64 i = 0; i < ShapeUtil::TupleElementCount(device_subshape); in WriteTupleIndexTablesAsync()
314 return WriteSingleTupleIndexTable(stream, elements, device_subshape, in WriteTupleIndexTablesAsync()
|