Dimension order represents how dimensions are laid out in memory, starting from the outer-most to the inner-most dimension. Thus, the conversion from strides is done by sorting the strides from larger to smaller since the dimension with the largest stride is the outer-most and t
(stride: Tuple[int])
| 52 | |
| 53 | |
| 54 | def dim_order_from_stride(stride: Tuple[int]) -> Tuple[bytes]: |
| 55 | """ |
| 56 | Dimension order represents how dimensions are laid out in memory, |
| 57 | starting from the outer-most to the inner-most dimension. |
| 58 | Thus, the conversion from strides is done by sorting the strides |
| 59 | from larger to smaller since the dimension with the largest stride |
| 60 | is the outer-most and the dimension with the smallest stride is the inner-most. |
| 61 | For example, tensor with sizes = (3, 5, 2) and strides = (5, 1, 15), implies |
| 62 | dimension order of (2, 0, 1). Dimension order of (2, 0, 1) can be obtained |
| 63 | by sorting strides from large to smaller. |
| 64 | |
| 65 | When strides do not convey dimension order unambiguously, dimension order |
| 66 | returned is dependent on stability of sort. In python same key elements are kept |
| 67 | in original order. Thus when strides = (4, 3, 1, 1) returned value is (0, 1, 2, 3) |
| 68 | Another example is: sizes = (1, 3, 1, 1) with strides = (3, 1, 3, 3), returned |
| 69 | value is (0, 2, 3, 1) |
| 70 | """ |
| 71 | from torch.fx.experimental.symbolic_shapes import ( |
| 72 | guard_or_false, |
| 73 | guard_size_oblivious, |
| 74 | ) |
| 75 | |
| 76 | for _, s in enumerate(stride): |
| 77 | if guard_or_false(s == 0): |
| 78 | raise ValueError("0 in strides is not supported for ExecuTorch.") |
| 79 | |
| 80 | class K(NamedTuple): |
| 81 | stride: int |
| 82 | |
| 83 | def __lt__(self, other): |
| 84 | return guard_size_oblivious(self.stride < other.stride) |
| 85 | |
| 86 | def __gt__(self, other): |
| 87 | return guard_size_oblivious(self.stride > other.stride) |
| 88 | |
| 89 | def __le__(self, other): |
| 90 | return guard_size_oblivious(self.stride <= other.stride) |
| 91 | |
| 92 | def __ge__(self, other): |
| 93 | return guard_size_oblivious(self.stride >= other.stride) |
| 94 | |
| 95 | def __eq__(self, other): |
| 96 | return guard_size_oblivious(self.stride == other.stride) |
| 97 | |
| 98 | sorted_dims = [ |
| 99 | i[0] for i in sorted(enumerate(stride), key=lambda x: K(x[1]), reverse=True) |
| 100 | ] |
| 101 | return tuple(typing.cast(Tuple[bytes], sorted_dims)) |
| 102 | |
| 103 | |
| 104 | def stride_from_dim_order(sizes: List[int], dim_order: List[int]) -> List[int]: |