(payload: list[FlattenedTensorMetadata], buffer: torch.Tensor)
| 33 | |
| 34 | |
| 35 | def _extract_weights(payload: list[FlattenedTensorMetadata], buffer: torch.Tensor) -> _WEIGHTS_TYPE: |
| 36 | assert buffer is not None |
| 37 | weights: _WEIGHTS_TYPE = [] |
| 38 | for item in payload: |
| 39 | shape = item["shape"] |
| 40 | if isinstance(shape, list | tuple): |
| 41 | shape = torch.Size(shape) |
| 42 | assert isinstance(shape, torch.Size) |
| 43 | dtype, offset = item["dtype"], item["offset"] |
| 44 | size = dtype.itemsize * shape.numel() |
| 45 | tensor = buffer[offset : offset + size].view(dtype=dtype).view(shape) |
| 46 | weights.append((item["name"], tensor)) |
| 47 | return weights |
| 48 | |
| 49 | |
| 50 | def update_weights_from_ipc( |
no outgoing calls
no test coverage detected