1515from executorch .exir .dialects .edge ._ops import EdgeOpOverload
1616from executorch .exir .error import ExportError , ExportErrorType
1717from executorch .exir .lowered_backend_module import LoweredBackendModule
18+ from executorch .exir .passes .dim_order_ops_registry import DimOrderOpsMap
1819from executorch .exir .verification .arg_validator import (
1920 EdgeOpArgValidator ,
2021 RunHigherOrderOperatorError ,
2122)
23+
2224from torch ._dispatch .python import enable_python_dispatcher
2325from torch ._export .utils import _detect_fake_mode_from_gm
2426
@@ -44,7 +46,7 @@ def _check_tensors_are_contiguous(gm: GraphModule) -> None:
4446
4547def _check_valid_dim_order_ops (op , use_dim_order ) -> None :
4648 if use_dim_order :
47- if op in ( torch . ops . aten . _to_copy . default ,) :
49+ if op in DimOrderOpsMap :
4850 raise SpecViolationError (f"{ op } should not be used in dim_order mode" )
4951 else : # not using dim_order
5052 if op .namespace in ("dim_order_ops" ,):
@@ -249,7 +251,7 @@ def check_valid_edge_op(self, op):
249251 )
250252 )
251253 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 )
253255 self .check_valid_aten_op (op ._op )
254256
255257 if isinstance (op , types .FunctionType ):
0 commit comments