[GitHub] szha closed pull request #12187: Make sure input symbol names are unique in control flow operators.

[GitBox <gi...@apache.org>]

szha closed pull request #12187: Make sure input symbol names are unique in control flow operators.
URL: https://github.com/apache/incubator-mxnet/pull/12187
 
 
   

This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:

As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):

diff --git a/python/mxnet/symbol/contrib.py b/python/mxnet/symbol/contrib.py
index 3c81dcf9ee0..b551422320c 100644
--- a/python/mxnet/symbol/contrib.py
+++ b/python/mxnet/symbol/contrib.py
@@ -135,6 +135,12 @@ def _regroup(args, fmt):
     return ret, args
 
 
+# We want to generate a unique name for input symbols to a control flow
+# operator. The names are generated on purpose differently from the symbols
+# cut from the graph.
+def _get_sym_uniq_name(sym):
+    return '{}-{}'.format(sym.name, sym.attr('_value_index'))
+
 def _get_graph_inputs(subg):
     num_handles = ctypes.c_int(0)
     handles = ctypes.POINTER(SymbolHandle)()
@@ -280,9 +286,9 @@ def foreach(body, data, init_states, name="foreach"):
     # with AttrScope and prune the nodes without the special attribute.
     name = _get_unique_subgraph_name(name)
     with AttrScope(__subgraph_name__=name):
-        in_eles = [symbol.var(sym.name) for sym in flatten_data]
+        in_eles = [symbol.var(_get_sym_uniq_name(sym)) for sym in flatten_data]
         in_eles, _ = _regroup(in_eles, data_fmt)
-        states = [symbol.var(s.name) for s in init_flatten_states]
+        states = [symbol.var(_get_sym_uniq_name(s)) for s in init_flatten_states]
         states, _ = _regroup(states, copy.deepcopy(init_state_fmt))
         sym_out, sym_states = body(in_eles, states)
 
@@ -310,12 +316,14 @@ def foreach(body, data, init_states, name="foreach"):
     gin_names = input_syms.keys()
     # This array contains the symbols for the inputs of foreach.
     # They are ordered according to the inputs of the subgraph.
-    state_names = [sym.name for sym in init_flatten_states]
-    data_names = [sym.name for sym in flatten_data]
+    state_names = [_get_sym_uniq_name(sym) for sym in init_flatten_states]
+    data_names = [_get_sym_uniq_name(sym) for sym in flatten_data]
     cut_var_map = {sym.list_outputs()[0]:sym for sym in cut_syms}
     cut_var_names = cut_var_map.keys()
 
     subg_input_names = g.list_inputs()
+    assert len(set(subg_input_names)) == len(subg_input_names), \
+            "The inputs of the subgraph don't have unique names: " + str(subg_input_names)
     # ordered_ins contains input symbols in the following order:
     # data_syms, state_syms, followed by cut_vars and vars in the closure.
     ordered_ins = [x for x in flatten_data]
@@ -483,7 +491,7 @@ def _create_subgraph(graph_vars, graph_func, subgraph_name):
             # create new variables with the same name,
             # them feed them to the given func
             graph_vars, var_fmt = _flatten(graph_vars, "while loop_vars")
-            new_graph_vars = [symbol.var(sym.name) for sym in graph_vars]
+            new_graph_vars = [symbol.var(_get_sym_uniq_name(sym)) for sym in graph_vars]
             new_graph_vars, _ = _regroup(new_graph_vars, var_fmt)
             outputs, final_state, out_fmt, var_fmt = graph_func(new_graph_vars)
             # first `num_out_data` elements belong to `outputs`
@@ -517,17 +525,20 @@ def _union_inputs(*graphs):
                                 # to a `loc`, where inputs[loc] = sym
         for graph in graphs:
             # some loop_vars are inputs to `graph`, some are not
-            name_to_loop_vars = {sym.name: sym for sym in flatten_loop_vars}
+            name_to_loop_vars = {_get_sym_uniq_name(sym): sym for sym in flatten_loop_vars}
             # other inputs to `graph` created by cut_graph
             name_to_cut_g_syms = {sym.list_outputs()[0]: sym for sym in _cut_subgraph(graph)}
             # input_syms: all inputs to the `graph`
             name_to_input_syms = {sym.name: sym for sym in _get_graph_inputs(graph)}
             # also we collect the mapping from var's name to var's loc in loop_vars
-            name_to_var_locs = {sym.name: i for i, sym in enumerate(flatten_loop_vars)}
+            name_to_var_locs = {_get_sym_uniq_name(sym): i for i, sym in enumerate(flatten_loop_vars)}
             # collect arguments for each subgraph
             input_locs = []                         # results from the second step
             var_locs = [-1] * len(flatten_loop_vars)        # results from the third step
-            for name in graph.list_inputs():
+            subg_input_names = graph.list_inputs()
+            assert len(set(subg_input_names)) == len(subg_input_names), \
+                    "The inputs of the subgraph don't have unique names: " + str(subg_input_names)
+            for name in subg_input_names:
                 assert name in name_to_input_syms   # it should obviously hold
                 # name -> sym
                 if name in name_to_loop_vars:
diff --git a/tests/python/unittest/test_contrib_control_flow.py b/tests/python/unittest/test_contrib_control_flow.py
index eadd63c9424..1c23c916197 100644
--- a/tests/python/unittest/test_contrib_control_flow.py
+++ b/tests/python/unittest/test_contrib_control_flow.py
@@ -1707,6 +1707,92 @@ def step2(data, states):
     assert_almost_equal(res1.asnumpy(), res2.asnumpy(), rtol=1e-3, atol=1e-3)
 
 
+@with_seed()
+def test_uniq_name():
+    class ForeachLayer1(gluon.HybridBlock):
+        def __init__(self, prefix=None, params=None):
+            super(ForeachLayer1, self).__init__(prefix=prefix, params=params)
+
+        def hybrid_forward(self, F, inputs, states):
+            def step1(data, states):
+                return data + 1, states
+            out1, states1 = F.contrib.foreach(step1, inputs, states)
+            # The input variables have the same symbol name.
+            out, states = F.contrib.foreach(step1, out1, states1)
+            return out
+
+    class ForeachLayer2(gluon.HybridBlock):
+        def __init__(self, prefix=None, params=None):
+            super(ForeachLayer2, self).__init__(prefix=prefix, params=params)
+
+        def hybrid_forward(self, F, inputs, states):
+            def step1(data, states):
+                return data + 1, states
+            out1, states1 = F.contrib.foreach(step1, inputs, states)
+            def step2(data, states):
+                return data, [states[0] + states1[0] + F.squeeze(out1.slice_axis(axis=0, begin=0, end=1))]
+            # The input variables have the same symbol names.
+            # The free variables have the same symbol names as the input variables.
+            out, states = F.contrib.foreach(step2, out1, states1)
+            return out
+
+    class WhileLayer1(gluon.HybridBlock):
+        def __init__(self, prefix=None, params=None):
+            super(WhileLayer1, self).__init__(prefix=prefix, params=params)
+
+        def hybrid_forward(self, F, inputs, states):
+            def cond(state1, state2):
+                s = F.squeeze(state1.slice_axis(axis=0, begin=0, end=1))
+                return s == s
+            def step(state1, state2):
+                return state1 + 1, [state1, state2]
+            states = [states[0], states[0] + 1]
+            out1, states1 = F.contrib.while_loop(cond, step, states, max_iterations=5)
+            # The input variables have the same symbol name.
+            out, states = F.contrib.while_loop(cond, step, states1, max_iterations=5)
+            return out
+
+    class WhileLayer2(gluon.HybridBlock):
+        def __init__(self, prefix=None, params=None):
+            super(WhileLayer2, self).__init__(prefix=prefix, params=params)
+
+        def hybrid_forward(self, F, inputs, states):
+            def cond(state1, state2):
+                s = F.squeeze(state1.slice_axis(axis=0, begin=0, end=1))
+                return s == s
+            def step1(state1, state2):
+                return state1 + 1, [state1, state2]
+            states = [states[0], states[0] + 1]
+            out1, states1 = F.contrib.while_loop(cond, step1, states, max_iterations=5)
+            def step2(state1, state2):
+                return state1 + 1, [state1 + states1[0], state2 + states1[1]]
+            # The input variables have the same symbol name.
+            out, states = F.contrib.while_loop(cond, step2, states1, max_iterations=5)
+            return out
+
+    TestLayers = [ForeachLayer1, ForeachLayer2,
+            WhileLayer1, WhileLayer2]
+
+    data = mx.nd.normal(loc=0, scale=1, shape=(2, 5))
+    states = mx.nd.normal(loc=0, scale=1, shape=(5))
+    for TestLayer in TestLayers:
+        layer = TestLayer()
+        layer.initialize(ctx=default_context())
+        res1 = layer(data, [states])
+
+        with mx.autograd.record():
+            res1 = layer(data, [states])
+
+        layer = TestLayer()
+        layer.initialize(ctx=default_context())
+        layer.hybridize()
+        res2 = layer(data, [states])
+
+        with mx.autograd.record():
+            res2 = layer(data, [states])
+        assert_almost_equal(res1.asnumpy(), res2.asnumpy(), rtol=0.001, atol=0.0001)
+
+
 @with_seed()
 def test_cut_subgraph_while_loop():
     class TestLayer(gluon.HybridBlock):


 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services