[MSC] add test for variable names with colons

Author: psunn

tests/python/contrib/test_msc/test_translate_relay.py

@@ -27,8 +27,11 @@
 import tvm.testing
 from tvm.relax.frontend.torch import from_fx
 from tvm.relay.frontend import from_pytorch
+from tvm import relay
+from tvm.ir.module import IRModule
 from tvm.contrib.msc.core.frontend import translate
 from tvm.contrib.msc.framework.tvm import codegen as tvm_codegen
+from tvm.contrib.msc.core import utils as msc_utils
 
 
 def _valid_target(target):
@@ -1057,5 +1060,37 @@ def forward(self, x, y):
     verify_model(Max(), [([256, 256], "float32"), ([256, 256], "float32")])
 
 
+def test_name_string_with_colon():
+    """test name string with colons,
+    e.g., TFLite default input name 'serving_default_input:0'
+    """
+
+    dtype = "float32"
+    x = relay.var("input_0:0", shape=(3, 5), dtype=dtype)
+    y = relay.var("input_1:0", shape=(3, 5), dtype=dtype)
+    z = relay.add(x, y)
+    func = relay.Function([x, y], z)
+    mod = IRModule()
+    mod["main"] = func
+
+    try:
+        graph, _ = translate.from_relay(mod)
+    except Exception as e:
+        raise RuntimeError(f"Translation from relay to graph failed: {e}")
+    inspect = graph.inspect()
+
+    expected = {
+        "inputs": [
+            {'name': 'input_0:0', 'shape': [3, 5], 'dtype': dtype, 'layout': ''},
+            {'name': 'input_1:0', 'shape': [3, 5], 'dtype': dtype, 'layout': ''}],
+        "outputs": [{'name': 'add', 'shape': [3, 5], 'dtype': dtype, 'layout': ''}],
+        "nodes": {'total': 3, 'input': 2, 'add': 1},
+        }
+
+    assert msc_utils.dict_equal(inspect, expected), "Inspect {} mismatch with expected {}".format(
+        inspect, expected
+    )
+
+
 if __name__ == "__main__":
     tvm.testing.main()