Update dim_order type

Differential Revision: D68041866

Pull Request resolved: #7610

exir/schema.py

@@ -64,7 +64,7 @@ class Tensor:
     scalar_type: ScalarType
     storage_offset: int
     sizes: List[int]
-    dim_order: List[bytes]
+    dim_order: List[int]
     requires_grad: bool
     layout: int
     data_buffer_idx: int

exir/tensor.py

@@ -76,7 +76,7 @@ def dim_order_from_stride(stride: Tuple[int]) -> Tuple[bytes]:
     return tuple(typing.cast(Tuple[bytes], sorted_dims))
 
 
-def stride_from_dim_order(sizes: List[int], dim_order: List[bytes]) -> List[int]:
+def stride_from_dim_order(sizes: List[int], dim_order: List[int]) -> List[int]:
     """
     Converts dim order to stride using sizes
     e.g. if sizes = (2, 3, 4) and dim_order = (0, 1, 2) then strides = (12, 4, 1)

exir/tests/common.py

@@ -5,8 +5,6 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-strict
-import typing
-from typing import List
 
 import torch
 
@@ -46,7 +44,7 @@ def get_test_program() -> Program:
                             scalar_type=ScalarType.FLOAT,
                             storage_offset=0,
                             sizes=[2, 2],
-                            dim_order=typing.cast(List[bytes], [0, 1]),
+                            dim_order=[0, 1],
                             requires_grad=False,
                             layout=0,
                             data_buffer_idx=0,