15
15
from executorch .exir .dialects .edge ._ops import EdgeOpOverload
16
16
from executorch .exir .error import ExportError , ExportErrorType
17
17
from executorch .exir .lowered_backend_module import LoweredBackendModule
18
+ from executorch .exir .passes .dim_order_ops_registry import DimOrderOpsMap
18
19
from executorch .exir .verification .arg_validator import (
19
20
EdgeOpArgValidator ,
20
21
RunHigherOrderOperatorError ,
21
22
)
23
+
22
24
from torch ._dispatch .python import enable_python_dispatcher
23
25
from torch ._export .utils import _detect_fake_mode_from_gm
24
26
@@ -44,7 +46,7 @@ def _check_tensors_are_contiguous(gm: GraphModule) -> None:
44
46
45
47
def _check_valid_dim_order_ops (op , use_dim_order ) -> None :
46
48
if use_dim_order :
47
- if op in ( torch . ops . aten . _to_copy . default ,) :
49
+ if op in DimOrderOpsMap :
48
50
raise SpecViolationError (f"{ op } should not be used in dim_order mode" )
49
51
else : # not using dim_order
50
52
if op .namespace in ("dim_order_ops" ,):
@@ -249,7 +251,7 @@ def check_valid_edge_op(self, op):
249
251
)
250
252
)
251
253
if isinstance (op , EdgeOpOverload ):
252
- _check_valid_dim_order_ops (op . _op , self .use_dim_order )
254
+ _check_valid_dim_order_ops (op , self .use_dim_order )
253
255
self .check_valid_aten_op (op ._op )
254
256
255
257
if isinstance (op , types .FunctionType ):
0 commit comments