Searched refs:device_subshape (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | generic_transfer_manager.cc | 122 [&](const Shape& device_subshape, const ShapeIndex& index) -> Status { in TransferLiteralToDeviceAsync() argument
124 if (device_subshape.IsArray()) { in TransferLiteralToDeviceAsync()
125 TF_RET_CHECK(GetByteSizeRequirement(device_subshape) == in TransferLiteralToDeviceAsync()
131 if (LayoutUtil::Equal(device_subshape.layout(), in TransferLiteralToDeviceAsync()
136 /*size=*/GetByteSizeRequirement(device_subshape), source, in TransferLiteralToDeviceAsync()
140 relayed_out_literal = subliteral.Relayout(device_subshape.layout(), in TransferLiteralToDeviceAsync()
145 /*size=*/GetByteSizeRequirement(device_subshape), source, in TransferLiteralToDeviceAsync()
|
| D | transfer_manager.cc | 245 [&](const Shape& device_subshape, const ShapeIndex& index) -> Status { in WriteTupleIndexTablesAsync() argument
246 if (device_subshape.IsTuple()) { in WriteTupleIndexTablesAsync()
248 TF_RET_CHECK(GetByteSizeRequirement(device_subshape) == in WriteTupleIndexTablesAsync()
253 for (int64 i = 0; i < ShapeUtil::TupleElementCount(device_subshape); in WriteTupleIndexTablesAsync()
259 return WriteSingleTupleIndexTable(stream, elements, device_subshape, in WriteTupleIndexTablesAsync()
|