[onnx_importer.py] Fix dim_value None not correctly processed and missing Float8E4M3FNUZType. (#4037)

As per title. Changes tested on SHARK-TestSuite's `alt_e2eshark`.

Author: giacs-epic

python/torch_mlir/extras/onnx_importer.py

@@ -55,6 +55,7 @@
     F16Type,
     F32Type,
     F64Type,
+    Float8E4M3FNUZType,
     Float8E4M3FNType,
     Float8E5M2FNUZType,
     Float8E5M2Type,
@@ -643,7 +644,7 @@ def get_list_element_type(self, tp: onnx.TypeProto) -> IrType:
         if tt.elem_type:
             element_type = self.tensor_element_type(tt.elem_type)
             dims = tuple(
-                (d.dim_value if not d.dim_param else None) for d in tt.shape.dim
+                (d.dim_value if d.HasField("dim_value") else None) for d in tt.shape.dim
             )
             shape_asm = ",".join("?" if d is None else str(d) for d in dims)
             return f"vtensor<[{shape_asm}],{element_type}>"
@@ -656,7 +657,7 @@ def get_optional_element_type(self, tp: onnx.TypeProto) -> IrType:
         if tt.elem_type:
             element_type = self.tensor_element_type(tt.elem_type)
             dims = tuple(
-                (d.dim_value if not d.dim_param else None) for d in tt.shape.dim
+                (d.dim_value if d.HasField("dim_value") else None) for d in tt.shape.dim
             )
             shape_asm = ",".join("?" if d is None else str(d) for d in dims)
             return f"vtensor<[{shape_asm}],{element_type}>"
@@ -707,13 +708,15 @@ def type_proto_to_type(self, tp: onnx.TypeProto) -> IrType:
 
         tt = tp.tensor_type
         if tt.elem_type:
-            if not tt.shape:
-                raise OnnxImportError(
-                    f"Unsupported Tensor type without shape (run shape inference?): {tp}"
-                )
             element_type = self.tensor_element_type(tt.elem_type)
             dims = tuple(
-                (d.dim_value if not d.dim_param else None) for d in tt.shape.dim
+                # NOTE: dynamic dimension can either be denoted by d.dim_param being set
+                #       (and d.dim_value consequently not set) or
+                #       by neither d.dim_value nor d.dim_param being set. Also note that
+                #       d.dim_value being 0 corresponds to the protobuf default when the field
+                #       is not set.
+                d.dim_value if d.HasField("dim_value") else None
+                for d in tt.shape.dim
             )
             return self.get_vtensor_type(dims, element_type)
 
@@ -1097,7 +1100,7 @@ def get_operator_function(
     onnx.TensorProto.DataType.COMPLEX128: lambda: ComplexType.get(F64Type.get()),
     onnx.TensorProto.DataType.BFLOAT16: lambda: BF16Type.get(),
     onnx.TensorProto.DataType.FLOAT8E4M3FN: lambda: Float8E4M3FNType.get(),
-    onnx.TensorProto.DataType.FLOAT8E4M3FNUZ: lambda: Float8E5M2FNUZType.get(),
+    onnx.TensorProto.DataType.FLOAT8E4M3FNUZ: lambda: Float8E4M3FNUZType.get(),
     onnx.TensorProto.DataType.FLOAT8E5M2: lambda: Float8E5M2Type.get(),
     onnx.TensorProto.DataType.FLOAT8E5M2FNUZ: lambda: Float8E5M2FNUZType.get(),
     onnx.TensorProto.DataType.STRING: lambda: "!torch.str",

python/torch_mlir/tools/import_onnx/__main__.py

@@ -86,15 +86,19 @@ def load_onnx_model(args: argparse.Namespace) -> onnx.ModelProto:
         raw_model = onnx.load(args.input_file, load_external_data=False)
         onnx.load_external_data_for_model(raw_model, str(args.data_dir))
 
+    raw_model_modified = False
+
     if args.opset_version:
         raw_model = onnx.version_converter.convert_version(
             raw_model, args.opset_version
         )
+        raw_model_modified = True
 
     if args.clear_domain:
         graph = raw_model.graph
         for n in graph.node:
             n.ClearField("domain")
+        raw_model_modified = True
 
     # Run the checker to test whether the file is above the threshold for
     # in-memory shape inference.  If not, go ahead and do the shape inference.
@@ -119,9 +123,15 @@ def load_onnx_model(args: argparse.Namespace) -> onnx.ModelProto:
 
     # Model is too big for in-memory inference: do file-based shape inference
     # to a temp file.
+    # First need to save as model when it has been changed (e.g. version conversion).
+    if raw_model_modified:
+        temp_raw_file = temp_dir / "raw.onnx"
+        onnx.save(raw_model, temp_raw_file, save_as_external_data=True)
     temp_inferred_file = temp_dir / "inferred.onnx"
     onnx.shape_inference.infer_shapes_path(
-        args.input_file, temp_inferred_file, data_prop=args.data_prop
+        temp_raw_file if raw_model_modified else args.input_file,
+        temp_inferred_file,
+        data_prop=args.data_prop,
     )
 
     # Sanity check the shape-inferred model to be sure we have a good model