| 36 | |
| 37 | |
| 38 | def contiguous_stride_from_shape(shape: torch.Size) -> Tuple[int]: |
| 39 | strides = [] |
| 40 | accum = 1 |
| 41 | for sz in reversed(shape): |
| 42 | strides.append(accum) |
| 43 | # For sizes[i] == 0, treat it as 1 to be consistent with core Pytorch |
| 44 | # This preserves the PT equivalent behavior for dims with 0 elements |
| 45 | if isinstance(sz, int): |
| 46 | if sz != 0: |
| 47 | accum *= sz |
| 48 | else: |
| 49 | # Unbacked symints may error on the != 0 check |
| 50 | accum *= sz |
| 51 | return tuple(reversed(strides)) |
| 52 | |
| 53 | |
| 54 | def dim_order_from_stride(stride: Tuple[int]) -> Tuple[bytes]: |