| 67 | ShapedBuffer::~ShapedBuffer() {} |
| 68 | |
| 69 | StatusOr<ShapedBuffer> ShapedBuffer::SubShapedBuffer( |
| 70 | const ShapeIndex& index) const { |
| 71 | TF_ASSIGN_OR_RETURN(const Shape* host_sub_shape, |
| 72 | ShapeUtil::TryGetSubshape(on_host_shape(), index)); |
| 73 | TF_ASSIGN_OR_RETURN(const Shape* device_sub_shape, |
| 74 | ShapeUtil::TryGetSubshape(on_device_shape(), index)); |
| 75 | ShapedBuffer sub_shaped_buffer(*host_sub_shape, *device_sub_shape, platform_, |
| 76 | device_ordinal_); |
| 77 | TF_ASSIGN_OR_RETURN(ShapeTree<se::DeviceMemoryBase> sub_buffers, |
| 78 | buffers_.SubShapeTree(index)); |
| 79 | sub_shaped_buffer.set_buffers(std::move(sub_buffers)); |
| 80 | return std::move(sub_shaped_buffer); |
| 81 | } |
| 82 | |
| 83 | void ShapedBuffer::clear() { |
| 84 | for (auto& pair : buffers_) { |