Update stablehlo.py

pytorch · Nov 29, 2023 · b60eabe · b60eabe
1 parent 8384841
commit b60eabe
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/torch_xla/stablehlo.py b/torch_xla/stablehlo.py
@@ -199,6 +199,14 @@ class StableHLOModelBundle:
   stablehlo_funcs: List[StableHLOFunc]
 
 
+@dataclass
+class StableHLOExportOptions:
+  include_human_readable_text: bool = True
+  override_tracing_arguments: Optional[Tuple[Any]] = None
+  override_tracing_kwargs: Optional[Mapping[str, Any]] = None
+  save_weights: bool = True
+
+
 class XLAExportInterpreter(torch.fx.Interpreter):
 
   def __init__(self, module, device):
@@ -456,14 +464,6 @@ def _load_program_bundle(stablehlo_dir: os.PathLike) -> StableHLOModelBundle:
       state_dict=state_dict)
 
 
-@dataclass
-class StableHLOExportOptions:
-  include_human_readable_text: bool = True
-  override_tracing_arguments: Optional[Tuple[Any]] = None
-  override_tracing_kwargs: Optional[Mapping[str, Any]] = None
-  save_weights: bool = True
-
-
 def save_as_stablehlo(exported_model: 'ExportedProgram',
                       stablehlo_dir: os.PathLike,
                       options: Optional[StableHLOExportOptions] = None):