| 171 | } |
| 172 | |
| 173 | ScopedShapedBuffer ScopedShapedBuffer::TakeSubTree(ShapeIndexView index) { |
| 174 | const xla::Shape& sub_on_host_shape = |
| 175 | xla::ShapeUtil::GetSubshape(on_host_shape(), {index}); |
| 176 | const xla::Shape& sub_on_device_shape = |
| 177 | xla::ShapeUtil::GetSubshape(on_device_shape(), {index}); |
| 178 | |
| 179 | ScopedShapedBuffer output(sub_on_host_shape, sub_on_device_shape, |
| 180 | memory_allocator(), device_ordinal()); |
| 181 | auto src_it = buffers().find(index); |
| 182 | auto dst_it = output.buffers().begin(); |
| 183 | while (dst_it != output.buffers().end()) { |
| 184 | dst_it->second = src_it->second; |
| 185 | src_it->second = tensorflow::se::DeviceMemoryBase(nullptr, 0); |
| 186 | ++src_it; |
| 187 | ++dst_it; |
| 188 | } |
| 189 | return output; |
| 190 | } |
| 191 | |
| 192 | } // namespace xla |