Tracks information about a tensor, such as format and data type.
| 42 | |
| 43 | @mod.export() |
| 44 | class TensorInfo: |
| 45 | """ |
| 46 | Tracks information about a tensor, such as format and data type. |
| 47 | """ |
| 48 | |
| 49 | @staticmethod |
| 50 | def from_trt(io_info): |
| 51 | """ |
| 52 | Creates a Polygraphy ``TensorInfo`` instance from a TensorRT ``IAlgorithmIOInfo``. |
| 53 | |
| 54 | Args: |
| 55 | io_info (trt.IAlgorithmIOInfo): The algorithm I/O information. |
| 56 | |
| 57 | Returns: |
| 58 | TensorInfo |
| 59 | """ |
| 60 | return TensorInfo( |
| 61 | io_info.tensor_format, |
| 62 | io_info.dtype, |
| 63 | tuple(io_info.strides), |
| 64 | # These fields were added in 8.6 |
| 65 | util.try_getattr(io_info, "vectorized_dim"), |
| 66 | util.try_getattr(io_info, "components_per_element"), |
| 67 | ) |
| 68 | |
| 69 | def __init__(self, tensor_format, dtype, strides, vectorized_dim, components_per_element): |
| 70 | """ |
| 71 | Args: |
| 72 | tensor_format (trt.TensorFormat): The tensor format. |
| 73 | dtype (trt.DataType): The data type. |
| 74 | strides (Sequence[int]): The strides. |
| 75 | vectorized_dim (int): The index of the vectorized dimensions. |
| 76 | components_per_element (int): The number of components per element. |
| 77 | """ |
| 78 | check_is_instance(tensor_format, trt.TensorFormat, "tensor_format") |
| 79 | check_is_instance(dtype, trt.DataType, "dtype") |
| 80 | check_is_instance(strides, Sequence, "strides") |
| 81 | if vectorized_dim is not None: |
| 82 | check_is_instance(vectorized_dim, int, "vectorized_dim") |
| 83 | if components_per_element is not None: |
| 84 | check_is_instance(components_per_element, int, "components_per_element") |
| 85 | |
| 86 | self.tensor_format = tensor_format |
| 87 | self.dtype = dtype |
| 88 | self.strides = tuple(strides) |
| 89 | self.vectorized_dim = vectorized_dim |
| 90 | self.components_per_element = components_per_element |
| 91 | |
| 92 | def __eq__(self, other): |
| 93 | return self.__dict__ == other.__dict__ |
| 94 | |
| 95 | def __repr__(self): |
| 96 | return f"TensorInfo({str(self.tensor_format)}, {str(self.dtype)}, {self.strides}, {self.vectorized_dim}, {self.components_per_element})" |
| 97 | |
| 98 | def __hash__(self): |
| 99 | return hash((self.tensor_format, self.dtype, self.strides, self.vectorized_dim, self.components_per_element)) |
| 100 | |
| 101 |
no outgoing calls