Convert constant and mutable tensors from a single byte-blob into a list of individual tensors. Args: constant_segment: SubsegmentOffset with the offsets of each tensor. segment_data: byte data containing the tensors and padding. Not modified. Returns: List[Buffer]
(
constant_segment: SubsegmentOffsets, segment_data: bytes
)
| 625 | |
| 626 | |
| 627 | def _restore_constant_segment( |
| 628 | constant_segment: SubsegmentOffsets, segment_data: bytes |
| 629 | ) -> List[Buffer]: |
| 630 | """Convert constant and mutable tensors from a single byte-blob into a list of individual tensors. |
| 631 | |
| 632 | Args: |
| 633 | constant_segment: SubsegmentOffset with the offsets of each tensor. |
| 634 | segment_data: byte data containing the tensors and padding. Not modified. |
| 635 | |
| 636 | Returns: |
| 637 | List[Buffer] containing each tensor in a separate object. |
| 638 | """ |
| 639 | buffers: List[Buffer] = [] |
| 640 | for i in range(len(constant_segment.offsets)): |
| 641 | start_offset = constant_segment.offsets[i] |
| 642 | # Note: this is the original end offset plus any padding between it and the next start offset |
| 643 | end_offset = ( |
| 644 | constant_segment.offsets[i + 1] |
| 645 | if i < len(constant_segment.offsets) - 1 |
| 646 | else len(segment_data) |
| 647 | ) |
| 648 | buffers.append(Buffer(storage=segment_data[start_offset:end_offset])) |
| 649 | return buffers |
| 650 | |
| 651 | |
| 652 | def _restore_named_data( |
no test coverage detected