def gen_intermediate_tuple(x):
     y1 = relay.add(x, relay.const(1, "float32"))
     y2 = relay.add(x, relay.const(1, "float32"))
     y3 = relay.add(x, relay.const(1, "float32"))
     concat = relay.concatenate((y1, y2, y3), axis=1)
     out = relay.add(concat, relay.const(1, "float32"))
     return out

 def before():
     c = relay.const(c_data)
     x = relay.var("x")
     y = relay.Tuple([x, c])
     z = relay.add(y[1], c)
     z = relay.add(z, y[0])
     return relay.Function([x], z)

 def before(x):
     concat = gen_consecutive_tuple(x)
     pooled = relay.nn.max_pool2d(concat, pool_size=(2, 2), strides=(2, 2), padding=(0, 0))
     out = relay.add(pooled, relay.const(1, "float32"))
     out2 = relay.add(out, relay.const(1, "float32"))
     out_tup = relay.Tuple((out, out2))
     return relay.Function(relay.ir_pass.free_vars(out_tup), out_tup)

 def expected():
     x = relay.var("x", shape=(1, 16))
     y = relay.nn.relu(x)
     y = relay.add(y, relay.const(1.0, "float32"))
     y = relay.add(y, y)
     f = relay.Function([x], y)
     return f

 def convnet():
     """Alternating layout of simple convnet (from image super-resolution).
     """
     bias1 = relay.var('bias1', shape=(64,))
     bias2 = relay.var('bias2', shape=(64,))
     bias3 = relay.var('bias3', shape=(64,))
     bias4 = relay.var('bias4', shape=(64,))
     weight1 = relay.var('weight1', shape=(64, 1, 5, 5))
     weight2 = relay.var('weight2', shape=(64, 64, 3, 3))
     weight3 = relay.var('weight3', shape=(64, 64, 3, 3))
     weight4 = relay.var('weight4', shape=(64, 64, 3, 3))
     data = relay.var("x", shape=(1, 1, 224, 224))
     n00 = relay.nn.conv2d(data, weight1, padding=[2, 2], kernel_size=[5, 5])
     n01 = relay.expand_dims(bias1, axis=1, num_newaxis=2)
     n02 = relay.add(n00, n01)
     n03 = relay.nn.relu(n02)
     n04 = relay.nn.conv2d(n03, weight2, padding=[1, 1], kernel_size=[3, 3])
     n05 = relay.expand_dims(bias2, axis=1, num_newaxis=2)
     n06 = relay.add(n04, n05)
     n07 = relay.nn.relu(n06)
     n08 = relay.nn.conv2d(n07, weight3, padding=[1, 1], kernel_size=[3, 3])
     n09 = relay.expand_dims(bias3, axis=1, num_newaxis=2)
     n10 = relay.add(n08, n09)
     n11 = relay.nn.relu(n10)
     n12 = relay.nn.conv2d(n11, weight4, padding=[1, 1], kernel_size=[3, 3])
     n13 = relay.expand_dims(bias4, axis=1, num_newaxis=2)
     n14 = relay.add(n12, n13)
     n15 = relay.reshape(n14, newshape=[1, 1, 3, 3, 224, 224])
     n16 = relay.transpose(n15, axes=[0, 1, 4, 2, 5, 3])
     net = relay.reshape(n16, newshape=[1, 1, 672, 672])
     args = relay.ir_pass.free_vars(net)
     return relay.Function(args, net)

 def before():
     x = relay.var("x", shape=(1, 16))
     y1 = relay.nn.relu(x)
     y2 = relay.nn.relu(x)
     y1 = relay.add(y1, relay.const(1.0, "float32"))
     y2 = relay.add(y2, relay.const(1.0, "float32"))
     y = relay.add(y1, y2)
     f = relay.Function([x], y)
     return f

def test_func():
    x = relay.var("x", shape=(3, 2))
    y = relay.var("y")
    one = relay.const(10e10, dtype="float32")
    z = relay.add(x, one)
    z = relay.add(z, z)
    f = relay.Function([x, y], z)
    show(z.astext())
    show(f.astext())

def test_env():
    x = relay.var("x", "float32")
    y = relay.var("y", "float32")
    z = relay.add(x, y)
    z = relay.add(z, z)
    f = relay.Function([x, y], z)
    env = relay.Module()
    env["myf"] = f
    text = env.astext()
    assert "def @myf" in text
    assert "%1 = add(%0, %0) # ty=float32" in text
    show(env.astext(annotate=lambda x: str(x.checked_type.dtype)))
    show(text)

def test_bind_params():
    x = relay.var("x")
    y = relay.var("y")
    z = relay.add(x, y)
    f = relay.Function([x, y], z)
    fbinded = relay.bind(f, {x : relay.const(1, "float32")})
    fexpected =relay.Function(
        [y],
        relay.add(relay.const(1, "float32"),  y))
    assert relay.ir_pass.alpha_equal(fbinded, fexpected)

    zbinded = relay.bind(z, {y: x})
    zexpected = relay.add(x, x)
    assert relay.ir_pass.alpha_equal(zbinded, zexpected)

def test_env():
    x = relay.var("x", "float32")
    y = relay.var("y", "float32")
    z = relay.add(x, y)
    z = relay.add(z, z)
    f = relay.Function([x, y], z)
    env = relay.Module()
    env["myf"] = f
    text = env.astext()
    assert "def @myf" in text
    assert "def @myf" in str(env)
    assert "add(%0, %0) /* ty=float32 */" in text
    assert "add(%0, %0) /* ty=float32 */" in str(env)
    show(env.astext(annotate=lambda x: str(x.checked_type.dtype) if type(x) == relay.Call else ""))
    show(text)

    def annotated():
        conv2d_1 = relay.nn.conv2d(
            data1,
            weight,
            channels=64,
            kernel_size=(3, 3),
            padding=(1, 1))
        _conv2d_1 = relay.annotation.on_device(conv2d_1, dev2)
        conv2d_2 = relay.nn.conv2d(
            data2,
            weight,
            channels=64,
            kernel_size=(3, 3),
            padding=(1, 1))
        _conv2d_2 = relay.annotation.on_device(conv2d_2, dev2)
        add = relay.add(conv2d_1, conv2d_2)
        _add = relay.annotation.on_device(add, dev1)
        conv2d_3 = relay.nn.conv2d(
            add,
            weight,
            channels=64,
            kernel_size=(3, 3),
            padding=(1, 1))
        _conv2d_3 = relay.annotation.on_device(conv2d_3, dev2)

        func = relay.Function([data1, data2, weight],
                              relay.Tuple(tvm.convert([_conv2d_1, _conv2d_2,
                                                       _conv2d_3, _add,
                                                       conv2d_3])))
        func = relay.ir_pass.infer_type(func)
        func = relay.ir_pass.rewrite_annotated_ops(func,
                                                   tvm.context(3).device_type)
        func = relay.ir_pass.infer_type(func)
        return relay.Function(relay.ir_pass.free_vars(func.body[4]),
                              func.body[4])

    def before(x, conv_weight, out_bias, out_scale, channels):
        args = [x, conv_weight, out_bias]
        y0 = relay.nn.conv2d(x, conv_weight,
                             channels=channels,
                             kernel_size=(3, 3),
                             padding=(1, 1))
        y0 = relay.multiply(y0, out_scale)
        y0 = relay.nn.relu(y0)

        y1 = relay.nn.conv2d(y0, conv_weight,
                             channels=channels,
                             kernel_size=(3, 3),
                             padding=(1, 1))
        y1 = relay.multiply(y1, out_scale)
        y1 = relay.nn.relu(y1)

        y2 = relay.nn.conv2d(y0, conv_weight,
                             channels=channels,
                             kernel_size=(3, 3),
                             padding=(1, 1))
        y2 = relay.multiply(y2, out_scale)
        y2 = relay.nn.relu(y2)

        y = relay.add(y1, y2)
        return relay.Function(args, y)

 def expected():
     sb = relay.ScopeBuilder()
     x = relay.var("x")
     c_folded = (c_data + c_data)
     t3 = sb.let("t3", relay.add(relay.const(c_folded), x))
     sb.ret(t3)
     return relay.Function([x], sb.get())

 def expected(x, conv_weight, out_bias, out_scale, channels):
     # use a fixed order of args so alpha equal check can pass
     args = [x, conv_weight, out_bias]
     def fold_conv_weight():
         squeezed_scale = relay.squeeze(out_scale, axis=[1,2])
         return  relay.multiply(
             conv_weight ,
             relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
     y0 = relay.nn.conv2d(x, fold_conv_weight(),
                         channels=channels,
                         kernel_size=(3, 3),
                         padding=(1, 1))
     y0 = relay.nn.relu(y0)
     y1 = relay.nn.conv2d(y0, fold_conv_weight(),
                         channels=channels,
                         kernel_size=(3, 3),
                         padding=(1, 1))
     y1 = relay.nn.relu(y1)
     y2 = relay.nn.conv2d(y0, fold_conv_weight(),
                         channels=channels,
                         kernel_size=(3, 3),
                         padding=(1, 1))
     y2 = relay.nn.relu(y2)
     y = relay.add(y1, y2)
     return relay.Function(args, y)

def test_depthwise_conv2d():
    batch_size = 1
    dshape = (batch_size, 64, 56, 56)
    weight_conv = relay.var("weight_depthwiseconv", shape=(64, 1, 3, 3))
    data1 = relay.var("data1", shape=dshape)
    data2 = relay.var("data2", shape=dshape)
    depthwise_conv2d_1 = relay.nn.conv2d(
        data1,
        weight_conv,
        kernel_size=(3, 3),
        padding=(1, 1),
        groups=64)
    depthwise_conv2d_2 = relay.nn.conv2d(
        data2,
        weight_conv,
        kernel_size=(3, 3),
        padding=(1, 1),
        groups=64)
    add = relay.add(depthwise_conv2d_1, depthwise_conv2d_2)
    func = relay.Function([data1, data2, weight_conv],
                            relay.Tuple(tvm.convert([depthwise_conv2d_1,
                                                    depthwise_conv2d_2,
                                                    add])))
    func = relay.ir_pass.infer_type(func)
    compute_count = relay.ir_pass.get_total_mac_number(func)
    assert compute_count == 2 * np.prod(dshape) * 3*3

 def before():
     x = relay.var("x", shape=(1, 64, 56, 56))
     weight = relay.var("weight")
     y = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1))
     y = relay.add(y, relay.const(1, "float32"))
     y = relay.Function(free_vars(y), y)
     return y

def test_plan_memory():
    # it is sufficient to cycle through two memories.

    x = relay.var("x", shape=(10,))
    y = relay.var("x", shape=(1,))
    y2 = relay.exp(y)
    z = relay.add(x, y2)
    z = relay.exp(z)
    z = relay.exp(z)
    z = relay.exp(z)
    z = relay.exp(z)
    z = relay.exp(z)
    func = relay.Function([x, y], z)
    func = relay.ir_pass.infer_type(func)
    func = relay.ir_pass.fuse_ops(func, opt_level=0)
    func = relay.ir_pass.infer_type(func)
    smap = relay.backend._backend.GraphPlanMemory(func)
    storage_ids = set()
    device_types = set()
    for k, v in smap.items():
        assert len(v) == 2
        for x in v[0]:
            storage_ids.add(x.value)
        for x in v[1]:
            device_types.add(x.value)

    # Current rule requires vars have unique storage id
    # because we don't do inplace, we will need another
    # two alternating temporary space.
    assert len(storage_ids) == 4
    assert len(device_types) == 1

    def expected():
        conv2d_1 = relay.nn.conv2d(
            data1,
            weight,
            channels=64,
            kernel_size=(3, 3),
            padding=(1, 1))
        device_copy1 = relay.device_copy(conv2d_1, dev2, dev1)
        conv2d_2 = relay.nn.conv2d(
            data2,
            weight,
            channels=64,
            kernel_size=(3, 3),
            padding=(1, 1))
        device_copy2 = relay.device_copy(conv2d_2, dev2, dev1)
        add = relay.add(device_copy1, device_copy2)
        device_copy3 = relay.device_copy(add, dev1, dev2)
        conv2d_3 = relay.nn.conv2d(
            device_copy3,
            weight,
            channels=64,
            kernel_size=(3, 3),
            padding=(1, 1))

        func = relay.Function([data1, weight, data2], conv2d_3)
        return func

def test_binds():
    x = relay.var("x")
    y = relay.add(x, x)
    intrp = create_executor("debug")
    xx = np.ones((10, 20))
    res = intrp.evaluate(y, binds={x: xx}).asnumpy()
    tvm.testing.assert_allclose(xx + xx, res)

 def expected():
     x = relay.var("p", shape=(10, 20))
     y = relay.add(x, relay.const(1, "float32"))
     z = relay.exp(y)
     f1 = relay.Function([x], z)
     x = relay.var("x", shape=(10, 20))
     y = relay.Call(f1, [x])
     return relay.Function([x], y)