Skip to content

Commit 85b86b3

Browse files
authored
[onnx] Fix importer variable names to make mlir legal (#2690)
Some names for `onnx` identifiers are not legal in `mlir-ir`. Sanitize so that the generated `ir` is legal.
1 parent ccd469c commit 85b86b3

File tree

1 file changed

+7
-1
lines changed

1 file changed

+7
-1
lines changed

python/torch_mlir/extras/onnx_importer.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
from dataclasses import dataclass
3939

4040
import numpy as np
41+
import re
4142

4243
from ..ir import (
4344
ArrayAttr,
@@ -464,13 +465,18 @@ def type_proto_to_type(self, tp: onnx.TypeProto) -> IrType:
464465
# See TypeProto: sequence_type, map_type, optional_type, sparse_tensor_type.
465466
raise OnnxImportError(f"Unsupported ONNX TypeProto: {tp}")
466467

468+
def _sanitize_name(self, name):
469+
if not name.isidentifier():
470+
name = "_" + name
471+
return re.sub("[:/]", "_", name)
472+
467473
def tensor_proto_to_attr(self, tp: onnx.TensorProto) -> Attribute:
468474
tensor_type = self.tensor_proto_to_builtin_type(tp)
469475
if tp.HasField("raw_data"):
470476
# Conveniently, DenseResourceElementsAttr shares the raw data
471477
# format. We just give it maximum numeric alignment.
472478
return DenseResourceElementsAttr.get_from_buffer(
473-
tp.raw_data, tp.name, tensor_type, alignment=8
479+
tp.raw_data, self._sanitize_name(tp.name), tensor_type, alignment=8
474480
)
475481
else:
476482
# We have to do a data type specific instantiation from proto fields.

0 commit comments

Comments
 (0)