Copies the tensors from the provided list into a Cord and tracks the offsets of each tensor. Args: constant_buffer: list of Buffers from which to extract constants from. Not modified. tensor_alignment: Alignment in bytes. Each tensor in the cord will be padded to align
(
constant_buffer: List[Buffer],
tensor_alignment: Optional[int] = None,
)
| 347 | |
| 348 | |
| 349 | def _extract_constant_segment( |
| 350 | constant_buffer: List[Buffer], |
| 351 | tensor_alignment: Optional[int] = None, |
| 352 | ) -> Tuple[Cord, List[int]]: |
| 353 | """Copies the tensors from the provided list into a Cord and tracks the offsets |
| 354 | of each tensor. |
| 355 | |
| 356 | Args: |
| 357 | constant_buffer: list of Buffers from which to extract constants from. Not modified. |
| 358 | tensor_alignment: Alignment in bytes. Each tensor in the cord will be padded to align |
| 359 | with this value. Defaults to ALIGNMENT. |
| 360 | |
| 361 | Returns: |
| 362 | A tuple of (constant segment, list of offsets for each tensor in the segment) |
| 363 | """ |
| 364 | constant_segment_data: Cord = Cord() |
| 365 | constant_segment_offsets: List[int] = [] |
| 366 | current_offset: int = 0 |
| 367 | for i in range(len(constant_buffer)): |
| 368 | buffer = constant_buffer[i] |
| 369 | constant_segment_data.append(buffer.storage) |
| 370 | buffer_length = len(buffer.storage) |
| 371 | pad_length = ( |
| 372 | padding_required(buffer_length, tensor_alignment) |
| 373 | if tensor_alignment is not None |
| 374 | else 0 |
| 375 | ) |
| 376 | if i < len(constant_buffer) - 1: |
| 377 | constant_segment_data.append(b"\x00" * pad_length) |
| 378 | constant_segment_offsets.append(current_offset) |
| 379 | current_offset += buffer_length + pad_length |
| 380 | |
| 381 | return constant_segment_data, constant_segment_offsets |
| 382 | |
| 383 | |
| 384 | def _extract_named_data( |
no test coverage detected