Converts the normal torch tensor to a flatbuffer tensor.
(
data_buffer_idx: int,
allocation_info: Optional[schema.AllocationDetails],
spec: TensorSpec,
)
| 348 | |
| 349 | |
| 350 | def make_tensor_value( |
| 351 | data_buffer_idx: int, |
| 352 | allocation_info: Optional[schema.AllocationDetails], |
| 353 | spec: TensorSpec, |
| 354 | ) -> schema.Tensor: |
| 355 | """ |
| 356 | Converts the normal torch tensor to a flatbuffer tensor. |
| 357 | """ |
| 358 | |
| 359 | def to_list( |
| 360 | x: Union[torch.Size, int, List[int], Tuple[int]] |
| 361 | ) -> Union[List[int], List[torch.Size]]: |
| 362 | if isinstance(x, torch.Size) or isinstance(x, tuple): |
| 363 | return list(x) |
| 364 | elif isinstance(x, int): |
| 365 | return [x] |
| 366 | else: |
| 367 | return x |
| 368 | |
| 369 | tensor_size = to_list(spec.shape) |
| 370 | tensor_dim_order = to_list(spec.dim_order) |
| 371 | |
| 372 | extra_tensor_info = spec.extra_tensor_info |
| 373 | # Propagate device from TensorSpec into ExtraTensorInfo for serialization. |
| 374 | # Note: we don't propagate Device on CPU; if no device info will be noticed, |
| 375 | # tensor_parser will automatic treat it as CPU:0, to prevent pte size |
| 376 | # regression as much as possible. |
| 377 | if spec.device != schema.DeviceType.CPU: |
| 378 | if extra_tensor_info is None: |
| 379 | extra_tensor_info = schema.ExtraTensorInfo( |
| 380 | device_type=spec.device, |
| 381 | device_index=spec.device_index, |
| 382 | ) |
| 383 | else: |
| 384 | extra_tensor_info.device_type = spec.device |
| 385 | extra_tensor_info.device_index = spec.device_index |
| 386 | |
| 387 | flatbuffer_tensor = schema.Tensor( |
| 388 | scalar_type=scalar_type_enum(spec.scalar_type), |
| 389 | # The runtime currently only supports tensors with offsets of zero. |
| 390 | storage_offset=0, |
| 391 | sizes=tensor_size, |
| 392 | dim_order=tensor_dim_order, |
| 393 | requires_grad=spec.requires_grad, |
| 394 | data_buffer_idx=data_buffer_idx, |
| 395 | allocation_info=allocation_info, |
| 396 | layout=layout_enum(spec.layout), |
| 397 | shape_dynamism=spec.shape_dynamism, |
| 398 | extra_tensor_info=extra_tensor_info, |
| 399 | ) |
| 400 | return flatbuffer_tensor |
| 401 | |
| 402 | |
| 403 | def check_spec(tensor: torch.Tensor, spec: TensorSpec) -> None: |