def test_infer_type(): x = sym.Variable('x') y = sym.add(x, x, name='add1') y = sym.cast(y, dtype=1, name="cast1") g = graph.create(y) g = g.apply('InferType') jgraph = json.loads(g.apply('SaveJSON').json_attr('json')) jnodes = jgraph['nodes'] jnode_row_ptr = jgraph['node_row_ptr'] nindex = {n['name']: i for i, n in enumerate(jnodes)} assert g.json_attr('dtype')[jnode_row_ptr[nindex["cast1"]]] == 1 assert g.json_attr('dtype')[jnode_row_ptr[nindex["add1"]]] == 0
def test_infer_type(): x = sym.Variable('x', dtype=0) y = sym.add(x, x, name='add1') y = sym.cast(y, dtype=1, name="cast1") g = graph.create(y) g._set_json_attr("dtype_attr_key", "dtype") g = g.apply('InferType') jgraph = json.loads(g.apply('SaveJSON').json_attr('json')) jnodes = jgraph['nodes'] jnode_row_ptr = jgraph['node_row_ptr'] nindex = {n['name']: i for i, n in enumerate(jnodes)} assert g.json_attr('dtype')[jnode_row_ptr[nindex["cast1"]]] == 1 assert g.json_attr('dtype')[jnode_row_ptr[nindex["add1"]]] == 0
def test_place_device(): x = sym.Variable('x', device_group="stage1") y = sym.add(x, x, name='add1') y = sym.cast(y, dtype=1, name="cast1") z = sym.add(y, y, device_group="stage2", name="add2") z = sym.add(z, sym.exp(y, device_group="stage2"), name="add3") g = graph.create(z) g._set_json_attr("device_group_attr_key", "device_group") g._set_json_attr("device_assign_map", {"stage1": 0, "stage2" : 1}, "dict_str_int") g._set_json_attr("device_copy_op", "cross_device_copy") g = g.apply("PlaceDevice") jgraph = json.loads(g.apply('SaveJSON').json_attr('json')) jnodes = jgraph['nodes'] jnode_row_ptr = jgraph['node_row_ptr'] nindex = {n['name']: i for i, n in enumerate(jnodes)} assert g.json_attr('device')[jnode_row_ptr[nindex["add2"]]] == 1 assert g.json_attr('device')[jnode_row_ptr[nindex["add3"]]] == 1 assert g.json_attr('device')[jnode_row_ptr[nindex["cast1"]]] == 0
def test_place_device(): x = sym.Variable('x', device_group="stage1") y = sym.add(x, x, name='add1') y = sym.cast(y, dtype=1, name="cast1") z = sym.add(y, y, device_group="stage2", name="add2") z = sym.add(z, sym.exp(y, device_group="stage2"), name="add3") g = graph.create(z) g._set_json_attr("device_group_attr_key", "device_group") g._set_json_attr("device_assign_map", {"stage1": 0, "stage2" : 1}, "dict_str_int") g._set_json_attr("device_copy_op", "cross_device_copy") g = g.apply("PlaceDevice") jgraph = json.loads(g.apply('SaveJSON').json_attr('json')) jnodes = jgraph['nodes'] jnode_row_ptr = jgraph['node_row_ptr'] nindex = {n['name']: i for i, n in enumerate(jnodes)} assert g.json_attr('device')[jnode_row_ptr[nindex["add2"]]] == 1 assert g.json_attr('device')[jnode_row_ptr[nindex["add3"]]] == 1 assert g.json_attr('device')[jnode_row_ptr[nindex["cast1"]]] == 0