Function test_replace_opr

imperative/python/test/unit/utils/test_network.py:98–141

()

def test_replace_opr():

    a = Tensor([1, 2])
    b = Tensor([3, 4])

    @trace(symbolic=True, capture_as_const=True)
    def fwd(a, b):
        return (a + b) * 2

    fwd(a, b)
    orig_model = io.BytesIO()
    fwd.dump(
        orig_model, arg_names=["a", "b"], output_names="o", optimize_for_inference=False
    )
    orig_model.seek(0)

    graph = Net.load(orig_model)
    vara = graph.var_filter.name("a").as_unique()
    varb = graph.var_filter.name("b").as_unique()

    out1 = F.mul(vara, varb)
    out1 = F.relu(out1)
    out1 += 2
    out1 *= 3
    out1 = graph.add_dep_oprs(out1)
    orig_opr = graph.opr_filter.has_input(vara).as_unique()

    new_opr = out1[0].owner
    repl_dict = {orig_opr: new_opr}
    graph.replace_oprs(repl_dict)

    var_out = orig_opr.outputs

    for idx, node in enumerate(var_out):
        assert node.owner is new_opr
        assert node.owner.outputs[idx] is node

    modified_model1 = io.BytesIO()
    graph.dump(modified_model1)
    modified_model1.seek(0)

    load_graph = GraphInference(modified_model1)
    out = load_graph.run(a, b)
    np.testing.assert_equal(out["o"], [30, 60])


def test_splice_network():