frontend/coreml/test_forward.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
import numpy as np

from coremltools.models.neural_network import NeuralNetworkBuilder
from coremltools.models import datatypes

import tvm
from tvm import te
from tvm.contrib import graph_runtime
from tvm import topi
import tvm.topi.testing
from tvm import relay
from tvm.topi.testing import conv2d_nchw_python

import coremltools as cm
import model_zoo
import tvm.testing


def get_tvm_output(
    func, x, params, target, ctx, out_shape=(1, 1000), input_name="image", dtype="float32"
):
    with tvm.transform.PassContext(opt_level=3):
        lib = relay.build(func, target, params=params)
    m = graph_runtime.GraphModule(lib["default"](ctx))
    # set inputs
    m.set_input(input_name, tvm.nd.array(x.astype(dtype)))
    m.run()
    # get outputs
    out = m.get_output(0, tvm.nd.empty(out_shape, dtype))
    return out.asnumpy()


def run_model_checkonly(model_file, model_name="", input_name="image"):
    model = cm.models.MLModel(model_file)
    x = model_zoo.get_cat_image()
    shape_dict = {input_name: x.shape}
    # Some Relay passes change operators on the fly. Ensuring that we generate
    # new graph for each target.
    for target, ctx in tvm.testing.enabled_targets():
        mod, params = relay.frontend.from_coreml(model, shape_dict)
        tvm_output = get_tvm_output(mod["main"], x, params, target, ctx)
        print(target, ctx, model_name, "prediction id: ", np.argmax(tvm_output.flat))


@tvm.testing.uses_gpu
def test_mobilenet_checkonly():
    model_file = model_zoo.get_mobilenet()
    run_model_checkonly(model_file, "mobilenet")


@tvm.testing.uses_gpu
def test_resnet50_checkonly():
    model_file = model_zoo.get_resnet50()
    run_model_checkonly(model_file, "resnet50")


def run_tvm_graph(
    coreml_model, target, ctx, input_data, input_name, output_shape, output_dtype="float32"
):
    """ Generic function to compile on relay and execute on tvm """
    if isinstance(input_data, list):
        shape_dict = {}
        dtype_dict = {}
        for i, e in enumerate(input_name):
            shape_dict[e] = input_data[i].shape
            dtype_dict[e] = input_data[i].dtype
    else:
        shape_dict = {input_name: input_data.shape}
        dtype_dict = {input_name: input_data.dtype}

    mod, params = relay.frontend.from_coreml(coreml_model, shape_dict)
    with tvm.transform.PassContext(opt_level=3):
        lib = relay.build(mod, target, params=params)

    from tvm.contrib import graph_runtime

    m = graph_runtime.GraphModule(lib["default"](ctx))
    # set inputs
    if isinstance(input_data, list):
        for i, e in enumerate(input_name):
            m.set_input(e, tvm.nd.array(input_data[i].astype(input_data[i].dtype)))
    else:
        m.set_input(input_name, tvm.nd.array(input_data.astype(input_data.dtype)))

    # execute
    m.run()
    # get outputs
    if isinstance(output_shape, list) and isinstance(output_dtype, list):
        tvm_output_list = []
        for i, s in enumerate(output_shape):
            tvm_output = m.get_output(i, tvm.nd.empty((s), output_dtype[i]))
            tvm_output_list.append(tvm_output.asnumpy())
        return tvm_output_list
    else:
        if not output_shape:
            tvm_output = m.get_output(0)
        else:
            tvm_output = m.get_output(0, tvm.nd.empty((output_shape), output_dtype))
        return tvm_output.asnumpy()


def verify_AddLayerParams(input_dim, alpha=2):
    dtype = "float32"

    a_np1 = np.random.uniform(size=input_dim).astype(dtype)
    a_np2 = np.random.uniform(size=input_dim).astype(dtype)

    b_np = np.add(a_np1, a_np2) + alpha
    inputs = [("input1", datatypes.Array(*input_dim)), ("input2", datatypes.Array(*input_dim))]
    output = [("output", datatypes.Array(*b_np.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_elementwise(
        name="Add", alpha=alpha, input_names=["input1", "input2"], output_name="output", mode="ADD"
    )
    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(
            model, target, ctx, [a_np1, a_np2], ["input1", "input2"], b_np.shape, dtype
        )
        tvm.testing.assert_allclose(out, b_np, rtol=1e-5)


@tvm.testing.uses_gpu
def test_forward_AddLayerParams():
    verify_AddLayerParams((1, 2, 2), 0)
    verify_AddLayerParams((1, 2, 2), 1)
    verify_AddLayerParams((1, 3, 3), 2)


def verify_MultiplyLayerParams(input_dim, alpha):
    dtype = "float32"

    a_np1 = np.random.uniform(size=input_dim).astype(dtype)
    a_np2 = np.random.uniform(size=input_dim).astype(dtype)

    b_np = np.multiply(a_np1, a_np2) * alpha
    inputs = [("input1", datatypes.Array(*input_dim)), ("input2", datatypes.Array(*input_dim))]
    output = [("output", datatypes.Array(*b_np.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_elementwise(
        name="Mul",
        alpha=alpha,
        input_names=["input1", "input2"],
        output_name="output",
        mode="MULTIPLY",
    )
    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(
            model, target, ctx, [a_np1, a_np2], ["input1", "input2"], b_np.shape, dtype
        )
        tvm.testing.assert_allclose(out, b_np, rtol=1e-5)


@tvm.testing.uses_gpu
def test_forward_MultiplyLayerParams():
    verify_MultiplyLayerParams((1, 2, 2), 0)
    verify_MultiplyLayerParams((1, 2, 2), 1)
    verify_MultiplyLayerParams((1, 3, 3), 2)


def verify_ConcatLayerParams(input1_dim, input2_dim):
    dtype = "float32"

    a_np1 = np.random.uniform(size=input1_dim).astype(dtype)
    a_np2 = np.random.uniform(size=input2_dim).astype(dtype)

    b_np = np.concatenate((a_np1, a_np2), axis=1)
    inputs = [("input1", datatypes.Array(*input1_dim)), ("input2", datatypes.Array(*input2_dim))]
    output = [("output", datatypes.Array(*b_np.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_elementwise(
        name="Concate", input_names=["input1", "input2"], output_name="output", mode="CONCAT"
    )
    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(
            model, target, ctx, [a_np1, a_np2], ["input1", "input2"], b_np.shape, dtype
        )
        tvm.testing.assert_allclose(out, b_np, rtol=1e-5)


@tvm.testing.uses_gpu
def test_forward_ConcatLayerParams():
    verify_ConcatLayerParams((1, 1, 2, 2), (1, 2, 2, 2))
    verify_ConcatLayerParams((1, 2, 4, 4), (1, 3, 4, 4))


def verify_UpsampleLayerParams(input_dim, scale, mode):
    dtype = "float32"

    a_np = np.full(input_dim, 1, dtype=dtype)
    if mode == "NN":
        b_np = tvm.topi.testing.upsampling_python(a_np, (scale, scale))
    else:
        new_h = input_dim[2] * scale
        new_w = input_dim[3] * scale
        b_np = tvm.topi.testing.bilinear_resize_python(a_np, (new_h, new_w), "NCHW")

    input = [("input", datatypes.Array(*input_dim))]
    output = [("output", datatypes.Array(*b_np.shape))]
    builder = NeuralNetworkBuilder(input, output)
    builder.add_upsample(
        name="Upsample",
        scaling_factor_h=scale,
        scaling_factor_w=scale,
        mode=mode,
        input_name="input",
        output_name="output",
    )

    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(model, target, ctx, a_np, "input", b_np.shape, dtype)
        tvm.testing.assert_allclose(out, b_np, rtol=1e-5)


@tvm.testing.uses_gpu
def test_forward_UpsampleLayerParams():
    verify_UpsampleLayerParams((1, 16, 32, 32), 2, "NN")
    verify_UpsampleLayerParams((1, 4, 6, 6), 3, "BILINEAR")


def verify_l2_normalize(input_dim, eps):
    dtype = "float32"

    a_np = np.random.uniform(size=input_dim).astype(dtype)
    b_np = tvm.topi.testing.l2_normalize_python(a_np, eps, 1)

    input = [("input", datatypes.Array(*input_dim))]
    output = [("output", datatypes.Array(*b_np.shape))]
    builder = NeuralNetworkBuilder(input, output)
    builder.add_l2_normalize(name="L2", epsilon=eps, input_name="input", output_name="output")

    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(model, target, ctx, a_np, "input", b_np.shape, dtype)
        tvm.testing.assert_allclose(out, b_np, rtol=1e-5)


@tvm.testing.uses_gpu
def test_forward_l2_normalize():
    verify_l2_normalize((1, 3, 20, 20), 0.001)


def verify_lrn(input_dim, size, bias, alpha, beta):
    dtype = "float32"
    axis = 1
    a_np = np.random.uniform(size=input_dim).astype(dtype)
    b_np = tvm.topi.testing.lrn_python(a_np, size, axis, bias, alpha, beta)

    input = [("input", datatypes.Array(*input_dim))]
    output = [("output", datatypes.Array(*b_np.shape))]
    builder = NeuralNetworkBuilder(input, output)
    builder.add_lrn(
        name="LRN",
        input_name="input",
        output_name="output",
        alpha=alpha,
        beta=beta,
        k=bias,
        local_size=size,
    )

    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(model, target, ctx, a_np, "input", b_np.shape, dtype)
        tvm.testing.assert_allclose(out, b_np, rtol=1e-5)


@tvm.testing.uses_gpu
def test_forward_lrn():
    verify_lrn((1, 3, 10, 20), 3, 1.0, 1.0, 0.5)


def verify_average(input_dim1, input_dim2, axis=0):
    dtype = "float32"

    a_np1 = np.random.uniform(size=input_dim1).astype(dtype)
    a_np2 = np.random.uniform(size=input_dim2).astype(dtype)

    b_np = np.mean((a_np1, a_np2), axis=axis)

    inputs = [("input1", datatypes.Array(*input_dim1)), ("input2", datatypes.Array(*input_dim2))]
    output = [("output", datatypes.Array(*b_np.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_elementwise(
        name="MEAN", input_names=["input1", "input2"], output_name="output", mode="AVE"
    )
    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(
            model, target, ctx, [a_np1, a_np2], ["input1", "input2"], b_np.shape, dtype
        )
        tvm.testing.assert_allclose(out, b_np, rtol=1e-5)


@tvm.testing.uses_gpu
def test_forward_average():
    verify_average((1, 3, 20, 20), (1, 3, 20, 20))
    verify_average((3, 20, 20), (1, 3, 20, 20))
    verify_average((20, 20), (1, 3, 20, 20))


def verify_max(input_dim):
    dtype = "float32"

    a_np1 = np.random.uniform(size=input_dim).astype(dtype)
    a_np2 = np.random.uniform(size=input_dim).astype(dtype)
    a_np3 = np.random.uniform(size=input_dim).astype(dtype)

    b_np = np.max((a_np1, a_np2, a_np3), axis=0)

    inputs = [
        ("input1", datatypes.Array(*input_dim)),
        ("input2", datatypes.Array(*input_dim)),
        ("input3", datatypes.Array(*input_dim)),
    ]
    output = [("output", datatypes.Array(*b_np.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_elementwise(
        name="Max", input_names=["input1", "input2", "input3"], output_name="output", mode="MAX"
    )
    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(
            model,
            target,
            ctx,
            [a_np1, a_np2, a_np3],
            ["input1", "input2", "input3"],
            b_np.shape,
            dtype,
        )
        tvm.testing.assert_allclose(out, b_np, rtol=1e-5)


@tvm.testing.uses_gpu
def test_forward_max():
    verify_max((1, 3, 20, 20))
    verify_max((20, 20))


def verify_min(input_dim):
    dtype = "float32"

    a_np1 = np.random.uniform(size=input_dim).astype(dtype)
    a_np2 = np.random.uniform(size=input_dim).astype(dtype)
    a_np3 = np.random.uniform(size=input_dim).astype(dtype)

    b_np = np.min((a_np1, a_np2, a_np3), axis=0)

    inputs = [
        ("input1", datatypes.Array(*input_dim)),
        ("input2", datatypes.Array(*input_dim)),
        ("input3", datatypes.Array(*input_dim)),
    ]
    output = [("output", datatypes.Array(*b_np.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_elementwise(
        name="Min", input_names=["input1", "input2", "input3"], output_name="output", mode="MIN"
    )
    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(
            model,
            target,
            ctx,
            [a_np1, a_np2, a_np3],
            ["input1", "input2", "input3"],
            b_np.shape,
            dtype,
        )
        tvm.testing.assert_allclose(out, b_np, rtol=1e-5)


@tvm.testing.uses_gpu
def test_forward_min():
    verify_min((1, 3, 20, 20))
    verify_min((20, 20))


def verify_unary_sqrt(input_dim):
    dtype = "float32"

    a_np = np.random.uniform(size=input_dim).astype(dtype)
    ref_val = np.sqrt(a_np)

    inputs = [("input", datatypes.Array(*input_dim))]
    output = [("output", datatypes.Array(*ref_val.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_unary(name="sqrt", input_name="input", output_name="output", mode="sqrt")

    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(model, target, ctx, [a_np], ["input"], ref_val.shape, dtype)
        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)


def verify_unary_rsqrt(input_dim, epsilon=0):
    dtype = "float32"

    a_np = np.random.uniform(size=input_dim).astype(dtype)
    ref_val = 1 / np.sqrt(a_np + epsilon)

    inputs = [("input", datatypes.Array(*input_dim))]
    output = [("output", datatypes.Array(*ref_val.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_unary(
        name="rsqrt", input_name="input", output_name="output", mode="rsqrt", epsilon=epsilon
    )

    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(model, target, ctx, [a_np], ["input"], ref_val.shape, dtype)
        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)


def verify_unary_inverse(input_dim, epsilon=0):
    dtype = "float32"

    a_np = np.random.uniform(size=input_dim).astype(dtype)
    ref_val = 1 / (a_np + epsilon)

    inputs = [("input", datatypes.Array(*input_dim))]
    output = [("output", datatypes.Array(*ref_val.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_unary(
        name="inverse", input_name="input", output_name="output", mode="inverse", epsilon=epsilon
    )

    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(model, target, ctx, [a_np], ["input"], ref_val.shape, dtype)
        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)


def verify_unary_power(input_dim, alpha):
    dtype = "float32"

    a_np = np.random.uniform(size=input_dim).astype(dtype)
    ref_val = np.power(a_np, alpha)

    inputs = [("input", datatypes.Array(*input_dim))]
    output = [("output", datatypes.Array(*ref_val.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_unary(
        name="power", input_name="input", output_name="output", mode="power", alpha=alpha
    )

    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(model, target, ctx, [a_np], ["input"], ref_val.shape, dtype)
        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)


def verify_unary_exp(input_dim):
    dtype = "float32"

    a_np = np.random.uniform(size=input_dim).astype(dtype)
    ref_val = np.exp(a_np)

    inputs = [("input", datatypes.Array(*input_dim))]
    output = [("output", datatypes.Array(*ref_val.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_unary(name="exp", input_name="input", output_name="output", mode="exp")

    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(model, target, ctx, [a_np], ["input"], ref_val.shape, dtype)
        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)


def verify_unary_log(input_dim):
    dtype = "float32"

    a_np = np.random.uniform(size=input_dim).astype(dtype)
    ref_val = np.log(a_np)

    inputs = [("input", datatypes.Array(*input_dim))]
    output = [("output", datatypes.Array(*ref_val.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_unary(name="log", input_name="input", output_name="output", mode="log")

    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(model, target, ctx, [a_np], ["input"], ref_val.shape, dtype)
        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)


def verify_unary_abs(input_dim):
    dtype = "float32"

    a_np = np.random.uniform(-100.0, 100.0, size=input_dim).astype(dtype)
    ref_val = np.abs(a_np)

    inputs = [("input", datatypes.Array(*input_dim))]
    output = [("output", datatypes.Array(*ref_val.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_unary(name="abs", input_name="input", output_name="output", mode="abs")

    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(model, target, ctx, [a_np], ["input"], ref_val.shape, dtype)
        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)


def verify_unary_threshold(input_dim, alpha):
    dtype = "float32"

    a_np = np.random.uniform(-100.0, 100.0, size=input_dim).astype(dtype)
    ref_val = np.maximum(a_np, alpha)

    inputs = [("input", datatypes.Array(*input_dim))]
    output = [("output", datatypes.Array(*ref_val.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_unary(
        name="threshold", input_name="input", output_name="output", mode="threshold", alpha=alpha
    )

    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(model, target, ctx, [a_np], ["input"], ref_val.shape, dtype)
        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)


@tvm.testing.uses_gpu
def test_forward_unary():
    verify_unary_sqrt((1, 3, 20, 20))
    verify_unary_rsqrt((1, 3, 20, 20))
    verify_unary_rsqrt((1, 3, 20, 20), epsilon=1e-6)
    verify_unary_inverse((1, 3, 20, 20))
    verify_unary_inverse((1, 3, 20, 20), epsilon=1e-6)
    verify_unary_power((1, 3, 20, 20), alpha=0.5)
    verify_unary_power((1, 3, 20, 20), alpha=4)
    verify_unary_exp((1, 3, 20, 20))
    verify_unary_log((1, 3, 20, 20))
    verify_unary_abs((1, 3, 20, 20))
    verify_unary_threshold((1, 3, 20, 20), alpha=-6.0)
    verify_unary_threshold((1, 3, 20, 20), alpha=5.0)


@tvm.testing.uses_gpu
def test_forward_reduce():
    from enum import Enum

    class ReduceAxis(Enum):
        CHW = 0
        HW = 1
        C = 2
        H = 3
        W = 4

    def _verify_reduce(input_dim, mode, axis, ref_func, dtype="float32"):
        print(input_dim, mode, axis)
        a_np = np.random.uniform(size=input_dim).astype(dtype)

        # translate to axis from coreml format
        if axis == ReduceAxis.CHW:
            np_axis = (-3, -2, -1)
        elif axis == ReduceAxis.HW:
            np_axis = (-2, -1)
        elif axis == ReduceAxis.C:
            np_axis = -3
        elif axis == ReduceAxis.H:
            np_axis = -2
        elif axis == ReduceAxis.W:
            np_axis = -1

        if ref_func == np.argmax:
            ref_val = np.expand_dims(ref_func(a_np, np_axis), np_axis).astype(dtype)
        else:
            ref_val = ref_func(a_np, np_axis, keepdims=True)

        inputs = [("input", datatypes.Array(*input_dim))]
        output = [("output", datatypes.Array(*ref_val.shape))]
        builder = NeuralNetworkBuilder(inputs, output)
        builder.add_reduce(
            name=mode, input_name="input", output_name="output", axis=axis.name, mode=mode
        )

        model = cm.models.MLModel(builder.spec)
        for target, ctx in tvm.testing.enabled_targets():
            out = run_tvm_graph(model, target, ctx, [a_np], ["input"], ref_val.shape, dtype)
            tvm.testing.assert_allclose(out, ref_val, rtol=1e-5, atol=1e-5)

    dshapes = [[10, 10], [1, 10, 10], [1, 3, 10, 10]]
    for dshape in dshapes:
        for axis in ReduceAxis:
            if len(dshape) < 3 and axis in [ReduceAxis.CHW, ReduceAxis.C]:
                # input must have rank at least 3
                continue
            _verify_reduce(dshape, "sum", axis, np.sum)
            _verify_reduce(dshape, "avg", axis, np.mean)
            _verify_reduce(dshape, "prod", axis, np.prod)
            _verify_reduce(dshape, "min", axis, np.min)
            _verify_reduce(dshape, "max", axis, np.max)
            if axis in [ReduceAxis.C, ReduceAxis.H, ReduceAxis.W]:
                # For mode ArgMax, axis must be [-1] or [-2] or [-3]
                _verify_reduce(dshape, "argmax", axis, np.argmax, dtype="int32")


def verify_reshape(input_dim, target_shape, mode):
    dtype = "float32"

    a_np = np.random.uniform(-100.0, 100.0, size=input_dim).astype(dtype)
    ref_val = np.reshape(a_np, target_shape)

    inputs = [("input", datatypes.Array(*input_dim))]
    output = [("output", datatypes.Array(*ref_val.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_reshape(
        name="reshape",
        input_name="input",
        output_name="output",
        target_shape=target_shape,
        mode=mode,
    )

    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(model, target, ctx, [a_np], ["input"], ref_val.shape, dtype)
        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)


def test_forward_reshape():
    for mode in [0, 1]:
        verify_reshape((20,), (1, 2, 2, 5), mode)
        verify_reshape((1, 3, 20, 20), (1, 12, 10, 10), mode)


def verify_split(input_dim, nOutputs):
    dtype = "float32"

    a_np = np.random.uniform(-100.0, 100.0, size=input_dim).astype(dtype)
    ref_val = np.split(a_np, nOutputs, axis=-3)

    inputs = [("input", datatypes.Array(*input_dim))]

    output_names = []
    outputs = []
    output_shapes = []
    for i, out in enumerate(ref_val):
        output_name = "output" + str(i)
        output_names = output_names + [output_name]
        outputs = outputs + [(output_name, datatypes.Array(*out.shape))]
        output_shapes = output_shapes + [out.shape]

    builder = NeuralNetworkBuilder(inputs, outputs)
    builder.add_split(name="split", input_name="input", output_names=output_names)

    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(
            model, target, ctx, [a_np], ["input"], output_shapes, [dtype] * len(output_shapes)
        )
        tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)


def test_forward_split():
    verify_split(
        (
            1,
            4,
            4,
            4,
        ),
        2,
    )
    verify_split(
        (
            1,
            3,
            30,
            20,
        ),
        3,
    )


def verify_image_scaler(input_dim, blue_bias=0.0, green_bias=0.0, red_bias=0.0, image_scale=1.0):
    dtype = "float32"
    a_np = np.random.uniform(size=input_dim).astype(dtype)
    # make sure it is valid image format CHW.
    assert len(a_np.shape) == 3 and a_np.shape[0] == 3
    b_np = np.zeros(a_np.shape, dtype=dtype)
    b_np[0, :, :] = image_scale * a_np[0, :, :] + blue_bias
    b_np[1, :, :] = image_scale * a_np[1, :, :] + green_bias
    b_np[2, :, :] = image_scale * a_np[2, :, :] + red_bias
    b_np = np.add(a_np, b_np)
    inputs = [("input1", datatypes.Array(*input_dim)), ("input2", datatypes.Array(*input_dim))]
    output = [("output", datatypes.Array(*b_np.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.set_pre_processing_parameters(
        image_input_names=["input1"],
        is_bgr=True,
        blue_bias=blue_bias,
        green_bias=green_bias,
        red_bias=red_bias,
        image_scale=image_scale,
    )
    # add one add layer to make CoreML model format valid
    # add layer has been tested before.
    builder.add_elementwise(
        name="add", input_names=["input1", "input2"], output_name="output", alpha=0, mode="ADD"
    )
    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(
            model, target, ctx, [a_np, a_np], ["input1", "input2"], b_np.shape, dtype
        )
        tvm.testing.assert_allclose(out, b_np, rtol=1e-5)


@tvm.testing.uses_gpu
def test_forward_image_scaler():
    verify_image_scaler((3, 224, 224), image_scale=0.17)
    verify_image_scaler(
        (3, 224, 224),
        blue_bias=-1.7669800519943237,
        green_bias=-1.985260009765625,
        red_bias=-2.102560043334961,
        image_scale=0.379,
    )


def verify_convolution(input_dim, filter, padding):
    dtype = "float32"
    N, C, H, W = input_dim
    OC, _, KH, KW = filter
    a_np = np.random.uniform(size=input_dim).astype(dtype)
    w_np = np.random.uniform(size=(OC, C, KH, KW)).astype(dtype)
    w_np_cm = np.transpose(w_np, axes=(2, 3, 1, 0))
    b_np = conv2d_nchw_python(a_np, w_np, [1, 1], padding)
    inputs = [("input1", datatypes.Array(C, H, W))]
    output = [("output", datatypes.Array(*b_np.shape))]
    builder = NeuralNetworkBuilder(inputs, output)
    builder.add_convolution(
        name="conv",
        kernel_channels=3,
        output_channels=OC,
        height=KH,
        width=KW,
        stride_height=1,
        stride_width=1,
        border_mode=padding.lower(),
        groups=1,
        W=w_np_cm,
        b=None,
        has_bias=False,
        is_deconv=False,
        input_name="input1",
        output_name="output",
    )
    model = cm.models.MLModel(builder.spec)
    for target, ctx in tvm.testing.enabled_targets():
        out = run_tvm_graph(model, target, ctx, [a_np], ["input1"], output_shape=None)
        tvm.testing.assert_allclose(out, b_np, rtol=1e-5)


@tvm.testing.uses_gpu
def test_forward_convolution():
    verify_convolution((1, 3, 224, 224), filter=(32, 3, 3, 3), padding="VALID")
    verify_convolution((1, 3, 224, 224), filter=(32, 3, 3, 3), padding="SAME")


if __name__ == "__main__":
    test_forward_AddLayerParams()
    test_forward_ConcatLayerParams()
    test_forward_MultiplyLayerParams()
    test_forward_UpsampleLayerParams()
    test_forward_l2_normalize()
    test_forward_lrn()
    test_forward_average()
    test_forward_max()
    test_forward_min()
    test_forward_unary()
    test_forward_reduce()
    test_forward_reshape()
    test_forward_split()
    test_mobilenet_checkonly()
    test_resnet50_checkonly()
    test_forward_image_scaler()
    test_forward_convolution()