From 1a6aa6a4472678d84c1e7314e1e989d2efff0bc7 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 4 Apr 2024 09:42:15 -0700 Subject: [PATCH] Include the ".onnxtext" extension in supported serialization format (#6051) Usage seen in https://github.com/onnx/onnx-mlir/blob/5ef91c1be51eb578e61d2ae3c05fa821297a8ee2/test/mlir/onnx/parse/fun_model_test.onnxtext. (https://github.com/onnx/onnx-mlir/pull/2262). I added the `.onnxtext` extension to be a known extension for the serializer. --------- Signed-off-by: Justin Chu --- onnx/__init__.py | 2 +- onnx/serialization.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/onnx/__init__.py b/onnx/__init__.py index 42a4ad023d2..e62c848501a 100644 --- a/onnx/__init__.py +++ b/onnx/__init__.py @@ -134,7 +134,7 @@ # Supported model formats that can be loaded from and saved to # The literals are formats with built-in support. But we also allow users to # register their own formats. So we allow str as well. -_SupportedFormat = Union[Literal["protobuf", "textproto"], str] +_SupportedFormat = Union[Literal["protobuf", "textproto", "onnxtxt", "json"], str] # Default serialization format _DEFAULT_FORMAT = "protobuf" diff --git a/onnx/serialization.py b/onnx/serialization.py index 1af780b761f..d0c57de1b15 100644 --- a/onnx/serialization.py +++ b/onnx/serialization.py @@ -171,7 +171,7 @@ class _TextualSerializer(ProtoSerializer): """Serialize and deserialize the ONNX textual representation.""" supported_format = "onnxtxt" - file_extensions = frozenset({".onnxtxt"}) + file_extensions = frozenset({".onnxtxt", ".onnxtext"}) def serialize_proto(self, proto: _Proto) -> bytes: text = onnx.printer.to_text(proto) # type: ignore[arg-type]