Refine fix to handle the case output is a TupleWrapper

Add a regression test guarding on original bug.
apache · Mar 13, 2019 · 7955a82 · 7955a82
1 parent ae2d046
commit 7955a82
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 6 deletions.
diff --git a/nnvm/python/nnvm/to_relay.py b/nnvm/python/nnvm/to_relay.py
@@ -441,12 +441,10 @@ def to_relay(graph, shape_dict, dtype_dict, params):
     graph = graph.apply(["InferShape", "InferType"])
     shape = graph.json_attr("shape")
     dtype = [graph_attr.TCODE_TO_DTYPE[di] for di in graph.json_attr("dtype")]
-    heads = [x[0] for x in json.loads(graph.json())['heads']]
 
     gidx = graph.index
     relay_map = {}
     fn_params = []
-    output_ids = []
 
     for nid, node in enumerate(gidx.nodes):
         children = []
@@ -468,9 +466,6 @@ def to_relay(graph, shape_dict, dtype_dict, params):
             fn_params.append(v)
             relay_map[nid] = v
         else:
-            if nid in heads:
-                output_ids.append(nid)
-
             if op_name in NNVM_OP_2_RELAY_OP:
                 str_attrs = StrAttrsDict(attrs)
                 call = NNVM_OP_2_RELAY_OP[op_name](children, str_attrs, odtype)
@@ -479,7 +474,14 @@ def to_relay(graph, shape_dict, dtype_dict, params):
                 raise Exception(
                     "nnvm.to_relay: unsupported operator: {0}".format(op_name))
 
-    outputs = [relay_map[nid] for nid in output_ids]
+    outputs = []
+    for nid, idx, _ in gidx.output_entries:
+        output = relay_map[nid]
+        if isinstance(output, expr.TupleWrapper):
+            outputs.append(output[idx])
+        else:
+            outputs.append(output)
+
     if len(outputs) == 1:
         body = outputs[0]
     else:

diff --git a/tests/python/frontend/nnvm_to_relay/test_forward.py b/tests/python/frontend/nnvm_to_relay/test_forward.py
@@ -72,6 +72,23 @@ def test_forward_dqn():
     verify_nnvm_to_relay(model, params, data_shape=(1, 4, 84, 84))
 
 
+def test_forward_split_concatenate():
+    shape = (2, 16)
+
+    tensor = nnvm.sym.Variable("data", shape=shape)
+
+    splited = nnvm.sym.split(tensor, indices_or_sections=2, axis=1)
+
+    concatenated = nnvm.sym.concatenate(*splited, axis=1)
+
+    params = {}
+
+    verify_nnvm_to_relay(splited[0], params, data_shape=shape)
+    verify_nnvm_to_relay(splited[1], params, data_shape=shape)
+    verify_nnvm_to_relay(splited, params, data_shape=shape)
+    verify_nnvm_to_relay(concatenated, params, data_shape=shape)
+
+
 if __name__ == '__main__':
     test_forward_mlp()
     test_forward_vgg()
@@ -80,3 +97,4 @@ def test_forward_dqn():
     test_forward_inception_v3()
     test_forward_densenet()
     test_forward_dqn()
+    test_forward_split_concatenate()