#  Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions
#  are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
#  EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
#  PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
#  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
#  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
#  OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

# Based on
#  https://github.com/NVIDIA/mxnet_to_onnx/blob/master/mx2onnx_converter/
# mx2onnx_converter_functions.py

# coding: utf-8
# pylint: disable=too-many-locals,no-else-return,too-many-lines
# pylint: disable=anomalous-backslash-in-string,eval-used
# pylint: disable=too-many-function-args
"""
Conversion Functions for common layers.
Add new functions here with a decorator.
"""

import re
import logging
import numpy as np
from .._export_onnx import MXNetGraph as mx_op
try:
    import onnx
except ImportError:
    onnx = None

OPSET_VERSION = 13

def parse_helper(attrs, attrs_name, alt_value=None):
    """Helper function to parse operator attributes in required format."""
    tuple_re = re.compile(r'\([0-9L|,| ]+\)')
    if not attrs:
        return alt_value
    attrs_str = None if attrs.get(attrs_name) is None else str(attrs.get(attrs_name))
    if attrs_str is None:
        return alt_value
    attrs_match = tuple_re.search(attrs_str)
    if attrs_match is not None:
        if attrs_match.span() == (0, len(attrs_str)):
            dims = eval(attrs_str)
            return dims
        else:
            raise AttributeError("Malformed %s dimensions: %s" % (attrs_name, str(attrs_str)))
    return alt_value

def transform_padding(pad_width):
    """Helper function to convert padding format for pad operator.
    """
    num_pad_values = len(pad_width)
    onnx_pad_width = [0]*num_pad_values

    start_index = 0
    # num_pad_values will always be multiple of 2
    end_index = int(num_pad_values/2)
    for idx in range(0, num_pad_values):
        if idx % 2 == 0:
            onnx_pad_width[start_index] = pad_width[idx]
            start_index += 1
        else:
            onnx_pad_width[end_index] = pad_width[idx]
            end_index += 1

    return onnx_pad_width


def convert_string_to_list(string_val):
    """Helper function to convert string to list.
     Used to convert shape attribute string to list format.
    """
    result_list = []

    list_string = string_val.split(',')
    for val in list_string:
        val = str(val.strip())
        val = val.replace("(", "")
        val = val.replace(")", "")
        val = val.replace("L", "")
        val = val.replace("[", "")
        val = val.replace("]", "")
        if val == "None":
            result_list.append(None)
        elif val != "":
            result_list.append(int(val))

    return result_list

def get_boolean_attribute_value(attrs, attr_name):
    """ Helper function to convert a string version
    of Boolean attributes to integer for ONNX.
    Takes attribute dictionary and attr_name as
    parameters.
    """
    return 1 if attrs.get(attr_name, 0) in ["True", "1"] else 0

def get_inputs(node, kwargs):
    """Helper function to get inputs"""
    name = node["name"]
    outputs_lookup = kwargs["outputs_lookup"]
    inputs = node["inputs"]
    attrs = node.get("attrs", {})
    input_nodes = []
    for ip in inputs:
        input_node_name = outputs_lookup[ip[0]][ip[1]].name
        input_nodes.append(input_node_name)

    return name, input_nodes, attrs

def get_input_dtypes(node, kwargs):
    outputs_lookup = kwargs['outputs_lookup']
    inputs = node['inputs']
    input_dtypes = []
    for ip in inputs:
        input_node_dtype = outputs_lookup[ip[0]][ip[1]].dtype
        input_dtypes.append(input_node_dtype)
    return input_dtypes

def create_basic_op_node(op_name, node, kwargs):
    """Helper function to create a basic operator
    node that doesn't contain op specific attrs"""
    name, input_nodes, _ = get_inputs(node, kwargs)

    node = onnx.helper.make_node(
        op_name,
        input_nodes,
        [name],
        name=name
    )
    return [node]

def create_const_scalar_node(input_name, value, kwargs):
    """Helper function to create a tensor value node and a
    initializer tensor node with constant value."""
    from onnx.helper import make_tensor
    initializer = kwargs["initializer"]
    dtype = value.dtype
    if dtype == 'float16':
        # when using float16, we must convert it to np.uint16 view first
        value = np.float16(value).view(np.uint16)
    input_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[dtype]
    tensor_node = make_tensor(input_name, input_type, (), ([value]))
    initializer.append(tensor_node)

def create_const_node(input_name, value, kwargs):
    """Helper function to create a tensor value node and a
    initializer tensor node with constant value."""
    from onnx.helper import make_tensor
    initializer = kwargs["initializer"]
    dtype = value.dtype
    if dtype == 'float16':
        # when using float16, we must convert it to np.uint16 view first
        value = np.float16(value).view(np.uint16)
    input_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[dtype]
    input_shape = value.shape
    tensor_node = make_tensor(input_name, input_type, input_shape, value)
    initializer.append(tensor_node)

def create_tensor(tensor_list, tensor_name, initializer, dtype='int64'):
    """Helper function to create a tensor value node and a
    initializer tensor node with constant value."""
    tensor_np = np.array(tensor_list, dtype=dtype)
    data_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[tensor_np.dtype]
    dims = np.shape(tensor_np)
    if dtype == np.float16:
        tensor_np = tensor_np.view(dtype=np.uint16)
    tensor = onnx.helper.make_tensor(
        name=tensor_name,
        data_type=data_type,
        dims=dims,
        vals=tensor_np.flatten().tolist(),
        raw=False
    )
    initializer.append(tensor)


def create_helper_trans_node(node_name, input_node):
    """create extra transpose node for dot operator"""
    trans_node = onnx.helper.make_node(
        'Transpose',
        inputs=[input_node],
        outputs=[node_name],
        name=node_name
    )
    return trans_node


def scalar_op_helper(node, op_name, **kwargs):
    """Helper function for scalar arithmetic operations"""
    from onnx import numpy_helper
    name, input_nodes, attrs = get_inputs(node, kwargs)
    input_dtypes = get_input_dtypes(node, kwargs)

    dtype = input_dtypes[0]
    dtype_t = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[dtype]

    scalar_value = np.array([attrs.get("scalar", 1)],
                            dtype=dtype)
    initializer = kwargs["initializer"]
    flag = True
    # If the input value is in initializer, just multiply with scalar input
    # and create a new initializer
    for i in initializer:
        if i.name == input_nodes[0]:
            if op_name == 'Mul':
                new_initializer = numpy_helper.to_array(i) * scalar_value[0]
            elif op_name == 'Sub':
                if name.startswith("_rminusscalar"):
                    new_initializer = scalar_value[0] - numpy_helper.to_array(i)
                else:
                    new_initializer = numpy_helper.to_array(i) - scalar_value[0]
            elif op_name == 'Add':
                new_initializer = numpy_helper.to_array(i) + scalar_value[0]
            elif op_name == 'Div':
                if name.startswith("_rdivscalar"):
                    new_initializer = scalar_value[0] / numpy_helper.to_array(i)
                else:
                    new_initializer = numpy_helper.to_array(i) / scalar_value[0]
            elif op_name == 'Pow':
                new_initializer = numpy_helper.to_array(i) ** scalar_value[0]
            flag = False
            break

    # else create a new tensor of the scalar value, add it in initializer
    if flag is True:
        dims = np.shape(scalar_value)

        scalar_op_name = "scalar_op" + str(kwargs["idx"])
        tensor_node = onnx.helper.make_tensor_value_info(scalar_op_name, dtype_t, dims)

        initializer.append(
            onnx.helper.make_tensor(
                name=scalar_op_name,
                data_type=dtype_t,
                dims=dims,
                vals=scalar_value,
                raw=False,
            )
        )

        mul_node = onnx.helper.make_node(
            op_name,
            [input_nodes[0], scalar_op_name],
            [name],
            name=name
        )

        return [tensor_node, mul_node]
    else:
        dtype_t = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[new_initializer.dtype]
        dims = np.shape(new_initializer)

        tensor_node = onnx.helper.make_tensor_value_info(name, dtype_t, dims)

        initializer.append(
            onnx.helper.make_tensor(
                name=name,
                data_type=dtype_t,
                dims=dims,
                vals=new_initializer.flatten(),
                raw=False,
            )
        )
        return [tensor_node]


    return create_basic_op_node('Shape', node, kwargs)


@mx_op.register("_contrib_arange_like", OPSET_VERSION)
def convert_arange_like(node, **kwargs):
    """Map MXNet's arange_like operator attributes to onnx's Range and Reshape operators.
    """
    from onnx.helper import make_node
    name, input_nodes, attrs = get_inputs(node, kwargs)
    input_dtypes = get_input_dtypes(node, kwargs)

    opset_version = kwargs['opset_version']
    if opset_version < 11:
        raise AttributeError("ONNX opset 11 or greater is required to export this operator")
    # use the same dtype as the that of the input node
    dtype = input_dtypes[0]
    dtype_t = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[dtype]
    axis = attrs.get('axis', 'None')
    start = attrs.get('start', 0.)
    step = attrs.get('step', 1.)
    repeat = int(attrs.get('repeat', 1))
    if repeat != 1:
        raise NotImplementedError("arange_like operator with repeat != 1 not yet implemented.")

    create_const_scalar_node(name+"_start", np.dtype(dtype).type(start), kwargs)
    create_const_scalar_node(name+"_step", np.dtype(dtype).type(step), kwargs)
    create_const_scalar_node(name+"_half_step", np.dtype(dtype).type(float(step)*0.5), kwargs)
    create_tensor([0], name+"_0", kwargs["initializer"], dtype='int64')
    nodes = []
    if axis == 'None':
        # output will be same shape as input
        nodes += [
            make_node("Shape", [input_nodes[0]], [name+"_shape0_out"]),
            make_node("ReduceProd", [name+"_shape0_out"], [name+"_redprod0_out"]),
            make_node("Squeeze", [name+"_redprod0_out", name+"_0"], [name+'_reshape0_out']),
            make_node("Cast", [name+"_reshape0_out"], [name+"_cast0_out"], to=dtype_t),
            make_node("Mul", [name+"_cast0_out", name+"_step"], [name+"_mul0_out"]),
            make_node("Add", [name+"_mul0_out", name+"_start"], [name+"_add1_out"]),
            make_node("Sub", [name+"_add1_out", name+"_half_step"], [name+"_sub0_out"]),
            make_node("Range", [name+"_start", name+"_sub0_out", name+"_step"], [name+"_range0_out"]),
            make_node("Reshape", [name+"_range0_out", name+"_shape0_out"], [name], name=name)
        ]
    else:
        # determine shape of axis
        create_tensor([int(axis)], name+"_axis_start", kwargs["initializer"], dtype='int64')
        create_tensor([int(axis)+1], name+"_axis_end", kwargs["initializer"], dtype='int64')
        nodes += [
            make_node("Shape", [input_nodes[0]], [name+"_shape0_out"]),
            make_node("Slice", [name+"_shape0_out", name+"_axis_start", name+"_axis_end"], [name+"_slice0_out"]),
            make_node("ReduceProd", [name+"_slice0_out"], [name+"_reprod0_out"]),
            make_node("Squeeze", [name+"_reprod0_out", name+"_0"], [name+"_reshape0_out"]),
            make_node("Cast", [name+"_reshape0_out"], [name+"_cast0_out"], to=dtype_t),
            make_node("Mul", [name+"_cast0_out", name+"_step"], [name+"_mul0_out"]),
            make_node("Add", [name+"_mul0_out", name+"_start"], [name+"_add1_out"]),
            make_node("Sub", [name+"_add1_out", name+"_half_step"], [name+"_sub0_out"]),
            make_node("Range", [name+"_start", name+"_sub0_out", name+"_step"], [name], name=name)
        ]

    return nodes


@mx_op.register("LayerNorm", OPSET_VERSION)
def convert_layer_norm(node, **kwargs):
    """Map MXNet's LayerNorm operator attributes to onnx operators.
    """
    from onnx.helper import make_node
    from onnx import TensorProto
    name, input_nodes, attrs = get_inputs(node, kwargs)
    input_dtypes = get_input_dtypes(node, kwargs)

    dtype = input_dtypes[0]

    axes = int(attrs.get('axis', -1))
    eps = attrs.get('eps', 9.99999975e-06)

    create_tensor([axes], name+"_axes", kwargs["initializer"])
    create_tensor([axes+1], name+"_axes+1", kwargs["initializer"])
    create_tensor([0], name+"_0", kwargs["initializer"], dtype='int64')
    create_const_scalar_node(name+'_0_s', np.int64(0), kwargs)
    create_const_scalar_node(name+'_1_s', np.int64(1), kwargs)
    create_const_scalar_node(name+"_2_s", np.int64(2).astype(dtype), kwargs)
    create_const_scalar_node(name+"_eps", np.float32(eps), kwargs)

    nodes = [
        make_node("ReduceMean", [input_nodes[0]], [name+"_rm0_out"], axes=[axes]),
        make_node("Sub", [input_nodes[0], name+"_rm0_out"], [name+"_sub0_out"]),
        make_node("Pow", [name+"_sub0_out", name+"_2_s"], [name+"_pow0_out"]),
        make_node("ReduceMean", [name+"_pow0_out"], [name+"_rm1_out"], axes=[axes]),
        make_node("Add", [name+"_rm1_out", name+"_eps"], [name+"_add0_out"]),
        make_node("Sqrt", [name+"_add0_out"], [name+"_sqrt0_out"]),
        make_node("Div", [name+"_sub0_out", name+"_sqrt0_out"], [name+"_div0_out"]),
    ]

    if axes == -1:
        nodes += [
            make_node("Mul", [name+"_div0_out", input_nodes[1]], [name+"_mul0_out"]),
            # make_node("Add", [name+"_mul0_out", input_nodes[2]], [name])
            # the Add operator triggers a weird NaN issue in onnxruntime
            # a workaround is to use Neg + Sub
            make_node('Neg', [input_nodes[2]], [name+'_neg']),
            make_node("Sub", [name+"_mul0_out", name+'_neg'], [name])
        ]
    else:
        nodes += [
            make_node("Shape", [input_nodes[0]], [name+"_shape0_out"]),
            make_node("Shape", [name+"_shape0_out"], [name+"_in_dim"]),
            make_node("Squeeze", [name+"_in_dim", name+"_0"], [name+"_in_dim_s"]),
            make_node("Range", [name+"_0_s", name+"_in_dim_s", name+"_1_s"], [name+"_range"]),
            make_node("Equal", [name+"_range", name+"_axes"], [name+"_equal"]),
            make_node("Cast", [name+"_equal"], [name+"_one_hot"], to=int(TensorProto.INT64)),
            make_node("Slice", [name+"_shape0_out", name+"_axes", name+"_axes+1"], [name+"_slice_out"]),
            make_node("Squeeze", [name+"_slice_out", name+"_0"], [name+"_slice_out_s"]),
            make_node("Sub", [name+"_slice_out_s", name+"_1_s"], [name+"_sub1_out"]),
            make_node("Mul", [name+"_one_hot", name+"_sub1_out"], [name+"_mul0_out"]),
            make_node("Add", [name+"_mul0_out", name+"_1_s"], [name+"_add1_out"]),
            make_node('Reshape', [input_nodes[1], name+"_add1_out"], [name+"gamma_exp"]),
            make_node('Reshape', [input_nodes[2], name+"_add1_out"], [name+"beta_exp"]),
            make_node('Expand', [name+"gamma_exp", name+"_shape0_out"], [name+"gamma_exp1"]),
            make_node('Expand', [name+"beta_exp", name+"_shape0_out"], [name+"beta_exp1"]),
            make_node("Mul", [name+"_div0_out", name+"gamma_exp1"], [name+"_mul1_out"]),
            make_node("Add", [name+"_mul1_out", name+"beta_exp1"], [name], name=name)
        ]

    return nodes


@mx_op.register("broadcast_axis", OPSET_VERSION)
def convert_broadcast_axis(node, **kwargs):
    """Map MXNet's broadcast_axis
    """
    from onnx.helper import make_node
    from onnx import TensorProto
    name, input_nodes, attrs = get_inputs(node, kwargs)

    axis = convert_string_to_list(attrs.get('axis', '()'))
    size = convert_string_to_list(attrs.get('size', '()'))
    assert len(axis) == len(size)

    shape_name = name+'_shape_0'

    create_tensor([0], name+'_0', kwargs["initializer"])
    create_tensor([1], name+'_1', kwargs["initializer"])
    create_const_scalar_node(name+'_0_s', np.int64(0), kwargs)
    create_const_scalar_node(name+'_1_s', np.int64(1), kwargs)

    nodes = [
        make_node('Shape', [input_nodes[0]], [shape_name]),
        make_node('Shape', [shape_name], [name+'_in_dim']),
        make_node('Squeeze', [name+'_in_dim', name+'_0'], [name+'_in_dim_s']),
        make_node('Range', [name+'_0_s', name+'_in_dim_s', name+'_1_s'], [name+'_range']),
    ]

    for i, axis in enumerate(axis):
        if axis not in (0, 1):
            create_tensor([axis], name+'_'+str(axis), kwargs["initializer"])
        create_tensor([size[i]-1], name+'_size_'+str(i), kwargs["initializer"])
        nodes += [
            make_node('Equal', [name+'_range', name+'_'+str(axis)], [name+'_equal_'+str(i)]),
            make_node('Cast', [name+'_equal_'+str(i)], [name+'_cast_'+str(i)], to=int(TensorProto.INT64)),
            make_node('Mul', [name+'_size_'+str(i), name+'_cast_'+str(i)], [name+'_mul_'+str(i)]),
            make_node('Add', [name+'_mul_'+str(i), name+'_1'], [name+'_add_'+str(i)]),
            make_node('Mul', [name+'_add_'+str(i), shape_name], [name+'_shape_'+str(i+1)])
        ]
        shape_name = name+'_shape_'+str(i+1)

    nodes += [
        make_node('Expand', [input_nodes[0], shape_name], [name], name=name)
    ]

    return nodes


@mx_op.register("SequenceMask", OPSET_VERSION)
def convert_sequencemask(node, **kwargs):
    """Map MXNet's SequenceMask operator
    """
    from onnx.helper import make_node
    from onnx import TensorProto

    name, input_nodes, attrs = get_inputs(node, kwargs)

    use_sequence_length = attrs.get('use_sequence_length', 'False')
    mask_val = float(attrs.get('value', '0'))
    axis = int(attrs.get('axis', '0'))

    if(use_sequence_length == 'False'):
        return [make_node('Identity', [input_nodes[0]], [name], name=name)]

    create_tensor([0], name+'_0', kwargs["initializer"])
    create_tensor([1], name+'_1', kwargs["initializer"])
    create_tensor([2], name+'_2', kwargs["initializer"])
    create_const_scalar_node(name+'_0_s', np.int64(0), kwargs)
    create_const_scalar_node(name+'_1_s', np.int64(1), kwargs)
    create_const_scalar_node(name+'_2_s', np.int64(2), kwargs)
    create_tensor([mask_val], name+'_mask_val', kwargs["initializer"], dtype='float32')

    nodes = [
        make_node('Shape', [input_nodes[0]], [name+'_in_shape']),
        make_node('Slice', [name+'_in_shape', name+'_0', name+'_1'], [name+'_slice_0']),
        make_node('Slice', [name+'_in_shape', name+'_1', name+'_2'], [name+'_slice_1']),
        make_node('Concat', [name+'_slice_0', name+'_1'], [name+'_shape_0'], axis=0),
        make_node('Shape', [name+'_in_shape'], [name+'_in_dim']),
        make_node('Squeeze', [name+'_in_dim', name+'_0'], [name+'_in_dim_s']),
        make_node('Range', [name+'_0_s', name+'_in_dim_s', name+'_1_s'], [name+'_range_0']),
        make_node('Less', [name+'_range_0', name+'_2'], [name+'_less_0']),
        make_node('Where', [name+'_less_0', name+'_in_shape', name+'_1'], [name+'_shape_1'])
    ]

    if(axis == 0):
        nodes += [
            make_node('Squeeze', [name+'_slice_0', name+'_0'], [name+'_max_len']),
            make_node('Range', [name+'_0_s', name+'_max_len', name+'_1_s'], [name+'_range_1']),
            make_node('Reshape', [name+'_range_1', name+'_shape_0'], [name+"_reshape_0"]),
            make_node('Cast', [input_nodes[1]], [name+'_cast'], to=int(TensorProto.INT64)),
            make_node('Less', [name+'_reshape_0', name+'_cast'], [name+'_less_1']),
            make_node('Reshape', [name+'_less_1', name+'_shape_1'], [name+"_reshape_1"]),
            make_node('Where', [name+'_reshape_1', input_nodes[0], name+'_mask_val'], [name], name=name),
        ]
    else:
        nodes += [
            make_node('Squeeze', [name+'_slice_1', name+'_0'], [name+'_max_len']),
            make_node('Range', [name+'_0_s', name+'_max_len', name+'_1_s'], [name+'_range_1']),
            make_node('Reshape', [input_nodes[1], name+'_shape_0'], [name+"_reshape_0"]),
            make_node('Cast', [name+"_reshape_0"], [name+'_cast'], to=int(TensorProto.INT64)),
            make_node('Less', [name+'_range_1', name+'_cast'], [name+'_less_1']),
            make_node('Reshape', [name+'_less_1', name+'_shape_1'], [name+"_reshape_1"]),
            make_node('Where', [name+'_reshape_1', input_nodes[0], name+'_mask_val'], [name], name=name),
        ]
    return nodes


@mx_op.register("expand_dims", OPSET_VERSION)
def convert_expand_dims(node, **kwargs):
    """Map MXNet's expand_dims operator attributes to onnx's Unsqueeze operator
    and return the created node.
    """
    name, input_nodes, attrs = get_inputs(node, kwargs)

    axis = int(attrs.get("axis"))
    create_tensor([axis], name+"_axis", kwargs["initializer"])
    input_nodes.append(name+"_axis")
    node = onnx.helper.make_node(
        "Unsqueeze",
        input_nodes,
        [name],
        name=name,
    )
    return [node]


@mx_op.register("stack", OPSET_VERSION)
def convert_stack(node, **kwargs):
    """Map MXNet's stack operator to onnx operators.
    """
    name, input_nodes, attrs = get_inputs(node, kwargs)
    axis = int(attrs.get("axis", 0))
    create_tensor([axis], name+"_axis", kwargs["initializer"])
    idx = 0
    nodes = []
    for input_node in input_nodes:
        nodes.append(onnx.helper.make_node(
            "Unsqueeze",
            inputs=[input_node, name+"_axis"],
            outputs=[name+"_unsqueeze"+str(idx)]
        ))
        idx += 1

    nodes.append(onnx.helper.make_node(
        "Concat",
        inputs=[name+"_unsqueeze"+str(i) for i in range(len(nodes))],
        outputs=[name],
        name=name,
        axis=axis
    ))
    return nodes


@mx_op.register("softmax", OPSET_VERSION)
def convert_softmax(node, **kwargs):
    """Map MXNet's softmax operator attributes to onnx's Softmax operator
    and return the created node.
    """
    from onnx.helper import make_node
    from onnx import TensorProto
    name, input_nodes, attrs = get_inputs(node, kwargs)
    input_dtypes = get_input_dtypes(node, kwargs)

    axis = int(attrs.get("axis", -1))
    temperature = str(attrs.get("temperature", 'None'))
    if temperature == 'None':
        temperature = 1.
    else:
        temperature = float(temperature)

    use_length = str(attrs.get("use_length", 'None'))
    dtype = input_dtypes[0]
    dtype_t = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[dtype]
    data = input_nodes[0]

    create_tensor([0], name+"_0", kwargs["initializer"])
    if axis == -1 and temperature == 1.:
        nodes = []
        if use_length == "True":
            # magic number, this is fp16 min
            create_tensor([-65500.0], name+"_mask_val", kwargs["initializer"], dtype=dtype)
            create_tensor([1], name+"_1", kwargs["initializer"])
            create_tensor([-1], name+"_-1", kwargs["initializer"])
            create_const_scalar_node(name+"_0_s", np.int64(0), kwargs)
            create_const_scalar_node(name+"_1_s", np.int64(1), kwargs)
            nodes += [
                make_node("Shape", [data], [name+"_shape"]),
                make_node("Shape", [name+"_shape"], [name+"_dim"]),
                make_node("Sub", [name+"_dim", name+"_1"], [name+"_dim_m1"]),
                make_node("Slice", [name+"_shape", name+"_dim_m1", name+"_dim"],
                          [name+"_dim_last_"]),
                make_node("Squeeze", [name+"_dim_last_", name+"_0"], [name+"_dim_last"]),
                make_node("Range", [name+"_0_s", name+"_dim_last", name+"_1_s"], [name+"_range"]),
                make_node("Cast", [input_nodes[1]], [name+"_len"], to=int(TensorProto.INT64)),
                make_node("Unsqueeze", [name+"_len", name+"_-1"], [name+"_len_unsqueezed"]),
                make_node("Less", [name+"_range", name+"_len_unsqueezed"], [name+"_less"]),
                make_node("Where", [name+'_less', data, name+"_mask_val"], [name+"_data_masked"])
            ]
            data = name+"_data_masked"

        nodes += [
            make_node("Softmax", [data], [name], axis=-1)
        ]

        return nodes

    create_tensor([axis], name+"_axes", kwargs["initializer"])
    create_tensor([temperature], name+"_tmp", kwargs["initializer"], dtype=dtype)
    nodes = [
        make_node("Div", [data, name+"_tmp"], [name+'_data']),
    ]
    if len(input_nodes) == 1:
        nodes += [
            make_node("Softmax", [name+'_data'], [name], axis=axis)
        ]
        return nodes
    elif use_length == "True":
        length = input_nodes[1]
        create_tensor([1], name+"_1", kwargs["initializer"])
        create_const_scalar_node(name+'_-1_s', np.int64(-1), kwargs)
        create_const_scalar_node(name+'_0_s', np.int64(0), kwargs)
        create_const_scalar_node(name+'_1_s', np.int64(1), kwargs)
        nodes += [
            # cast data type
            make_node("Cast", [length], [name+"_length"], to=int(TensorProto.INT64)),
            make_node("Cast", [name+"_0"], [name+"_0_itype"], to=dtype_t),
            make_node("Cast", [name+"_1"], [name+"_1_itype"], to=dtype_t),
            # softmax output
            make_node("Softmax", [name+'_data'], [name+"_softmax_out"], axis=axis),
            # update axis
            make_node("Shape", [data], [name+"_shape0_out"]),
            make_node("Shape", [name+"_shape0_out"], [name+"_in_dim"]),
            make_node("Add", [name+"_in_dim", name+"_axes"], [name+"_dim+axis"]),
            make_node("Less", [name+"_axes", name+"_0_s"], [name+"_less0_out"]),
            make_node("Where", [name+"_less0_out", name+"_dim+axis", name+"_axes"], [name+"_final_axis"]),
            # data mask
            make_node("Add", [name+"_final_axis", name+"_1_s"], [name+"_final_axis+1"]),
            make_node("Slice", [name+"_shape0_out", name+"_final_axis", name+"_final_axis+1"], [name+"_axis_dim"]),
            make_node("Squeeze", [name+"_axis_dim", name+"_0"], [name+"_axis_dim_s"]),
            make_node("Range", [name+"_0_s", name+"_axis_dim_s", name+"_1_s"], [name+"_range0_out"]),
            # one hot for axis
            make_node("Squeeze", [name+"_in_dim", name+"_0"], [name+"_in_dim_s"]),
            make_node("Range", [name+"_0_s", name+"_in_dim_s", name+"_1_s"], [name+"_range1_out"]),
            make_node("Equal", [name+"_range1_out", name+"_final_axis"], [name+"_equal_out"]),
            make_node("Cast", [name+"_equal_out"], [name+"_one_hot"], to=int(TensorProto.INT64)),
            # reshape data mask for less
            make_node("Sub", [name+"_axis_dim_s", name+"_1_s"], [name+"_sub0_out"]),
            make_node("Mul", [name+"_one_hot", name+"_sub0_out"], [name+"_mul0_out"]),
            make_node("Add", [name+"_mul0_out", name+"_1_s"], [name+"_add0_out"]),
            make_node('Reshape', [name+"_range0_out", name+"_add0_out"], [name+"_reshape0_out"]),
            # reshape length for less
            make_node("Mul", [name+"_one_hot", name+"_-1_s"], [name+"_mul1_out"]),
            make_node("Add", [name+"_mul1_out", name+"_1_s"], [name+"_add1_out"]),
            make_node("Sub", [name+"_shape0_out", name+"_1_s"], [name+"_sub1_out"]),
            make_node("Mul", [name+"_add1_out", name+"_sub1_out"], [name+"_mul2_out"]),
            make_node("Add", [name+"_mul2_out", name+"_1_s"], [name+"_add2_out"]),
            make_node('Reshape', [name+"_length", name+"_add2_out"], [name+"_reshape1_out"]),
            # mask output
            make_node("Less", [name+"_reshape0_out", name+"_reshape1_out"], [name+"_less_out"]),
            make_node("Cast", [name+"_less_out"], [name+"_mask"], to=dtype_t),
            make_node("Mul", [name+"_softmax_out", name+"_mask"], [name+"_mul3_out"]),
            make_node("ReduceSum", [name+"_mul3_out", name+"_axes"], [name+"_rsum1_out"], keepdims=1),
            make_node("Equal", [name+"_rsum1_out", name+"_0_itype"], [name+"_equal1_out"]),
            make_node("Where", [name+"_equal1_out", name+"_1_itype", name+"_rsum1_out"], [name+"_where_out"]),
            make_node("Div", [name+"_mul3_out", name+"_where_out"], [name], name=name)
        ]
        return nodes

    else:
        raise NotImplementedError("use_length must be true when both data and length are paased in.")


@mx_op.register("reverse", OPSET_VERSION)
def convert_reverse(node, **kwargs):
    """Map MXNet's reverse operator attributes to ONNX
    """
    from onnx.helper import make_node
    name, input_nodes, attrs = get_inputs(node, kwargs)

    axis = int(attrs.get('axis', 0))

    # Transpose takes perm as a parameter, so we must 'pad' the input to a known dim (8 here)
    perm = [i for i in range(8)]
    perm[0], perm[axis] = axis, 0

    create_tensor([8], name+'_8', kwargs['initializer'])
    create_tensor([0], name+'_0', kwargs['initializer'])
    create_tensor([1], name+'_1', kwargs['initializer'])
    create_tensor([-1], name+'_m1', kwargs['initializer'])
    create_tensor([axis], name+'_axis', kwargs['initializer'])
    create_tensor([axis+1], name+'_axis_p1', kwargs['initializer'])
    create_const_scalar_node(name+'_m1_s', np.int64(-1), kwargs)

    nodes = [
        make_node('Shape', [input_nodes[0]], [name+'_shape']),
        make_node('Shape', [name+'_shape'], [name+'_dim']),
        make_node('Sub', [name+'_8', name+'_dim'], [name+'_sub']),
        make_node('Concat', [name+'_0', name+'_sub'], [name+'_concat'], axis=0),
        make_node('Pad', [name+'_shape', name+'_concat', name+'_1'], [name+'_shape_8_dim']),
        make_node('Reshape', [input_nodes[0], name+'_shape_8_dim'], [name+'_data_8_dim']),
        make_node('Transpose', [name+'_data_8_dim'], [name+'_data_t'], perm=perm),
        make_node('Slice', [name+'_shape', name+'_axis', name+'_axis_p1'], [name+'_axis_len']),
        make_node('Sub', [name+'_axis_len', name+'_1'], [name+'_axis_len_m1']),
        make_node('Squeeze', [name+'_axis_len_m1', name+'_0'], [name+'_axis_len_m1_s']),
        make_node('Range', [name+'_axis_len_m1_s', name+'_m1_s', name+'_m1_s'], [name+'_indices']),
        make_node('Gather', [name+'_data_t', name+'_indices'], [name+'_gather']),
        make_node('Transpose', [name+'_gather'], [name+'_data_reversed'], perm=perm),
        make_node('Reshape', [name+'_data_reversed', name+'_shape'], [name], name=name)
    ]

    return nodes


@mx_op.register('repeat', OPSET_VERSION)
def convert_repeat(node, **kwargs):
    """Map MXNet's repeat operator attributes to onnx's Tile operator.
    """
    from onnx.helper import make_node
    from onnx import TensorProto
    name, input_nodes, attrs = get_inputs(node, kwargs)

    opset_version = kwargs['opset_version']
    if opset_version < 11:
        raise AttributeError('ONNX opset 11 or greater is required to export this operator')

    repeats = int(attrs.get('repeats', 1))
    axis = attrs.get('axis', 'None')

    if repeats <= 0:
        raise NotImplementedError('repeat operator does not support parameter repeats==0')

    nodes = []
    if axis == 'None':
        create_tensor([-1], name+'_-1', kwargs['initializer'])
        create_tensor([repeats], name+'_rep', kwargs['initializer'])
        create_tensor([1, repeats], name+'_repeats', kwargs['initializer'])
        nodes += [
            make_node('Shape', [input_nodes[0]], [name+'_shape']),
            make_node('ReduceProd', [name+'_shape'], [name+'_size']),
            make_node('Reshape', [input_nodes[0], name+'_size'], [name+'_flat']),
            make_node('Unsqueeze', [name+'_flat', name+'_-1'], [name+'_unsqueeze']),
            make_node('Tile', [name+'_unsqueeze', name+'_repeats'], [name+'_tile']),
            make_node('Mul', [name+'_size', name+'_rep'], [name+'_new_size']),
            make_node('Reshape', [name+'_tile', name+'_new_size'], [name], name=name)
        ]
    else:
        axis = int(axis)
        repeats -= 1
        create_tensor([repeats], name+'_repeats', kwargs['initializer'])
        create_tensor([1], name+'_1', kwargs['initializer'])
        create_tensor([0], name+'_0', kwargs['initializer'])
        create_tensor([axis], name+'_axis', kwargs['initializer'])
        create_const_scalar_node(name+"_0_s", np.int64(0), kwargs)
        create_const_scalar_node(name+"_1_s", np.int64(1), kwargs)
        nodes += [
            make_node('Shape', [input_nodes[0]], [name+'_shape']),
            make_node('Shape', [name+'_shape'], [name+'_dim']),
            make_node('Squeeze', [name+'_dim', name+'_0'], [name+'_dim_s']),
            make_node('Range', [name+'_0_s', name+'_dim_s', name+'_1_s'], [name+'_range'])
        ]
        if axis < 0:
            nodes += [
                make_node('Add', [name+'_axis', name+'_dim'], [name+'_true_axis']),
                make_node('Equal', [name+'_range', name+'_true_axis'], [name+'_one_hot'])
                ]
        else:
            nodes += [
                make_node('Equal', [name+'_range', name+'_axis'], [name+'_one_hot'])
                ]
        nodes += [
            make_node('Cast', [name+'_one_hot'], [name+'_one_hot_int'], to=int(TensorProto.INT64)),
            make_node('Mul', [name+'_repeats', name+'_one_hot_int'], [name+'_mul']),
            make_node('Add', [name+'_mul', name+'_1'], [name+'_add']),
            make_node('Concat', [name+'_1', name+'_add'], [name+'_repeats_tensor'], axis=0)
            ]
        if axis == -1:
            nodes += [
                make_node('Concat', [name+'_shape', name+'_1'], [name+'_unsqueeze_shape'], axis=0),
                make_node('Reshape', [input_nodes[0], name+'_unsqueeze_shape'],
                          [name+'_unsqueeze'])
                ]
        else:
            create_tensor([axis+1], name+'_axis+1', kwargs['initializer'])
            nodes += [
                make_node('Unsqueeze', [input_nodes[0], name+'_axis+1'], [name+'_unsqueeze'])
                ]
        nodes += [
            make_node('Tile', [name+'_unsqueeze', name+'_repeats_tensor'], [name+'_tile']),
            make_node('Mul', [name+'_shape', name+'_add'], [name+'_new_shape']),
            make_node('Reshape', [name+'_tile', name+'_new_shape'], [name], name=name)
            ]

    return nodes


@mx_op.register('_contrib_box_nms', OPSET_VERSION)
def convert_contrib_box_nms(node, **kwargs):
    """Map MXNet's _contrib_box_nms operator to ONNX
    """
    from onnx.helper import make_node
    name, input_nodes, attrs = get_inputs(node, kwargs)
    input_dtypes = get_input_dtypes(node, kwargs)

    dtype = input_dtypes[0]
    #dtype_t = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[dtype]

    opset_version = kwargs['opset_version']
    if opset_version < 11:
        raise AttributeError('ONNX opset 11 or greater is required to export this operator')

    overlap_thresh = float(attrs.get('overlap_thresh', '0.5'))
    valid_thresh = float(attrs.get('valid_thresh', '0'))
    topk = int(attrs.get('topk', '-1'))
    coord_start = int(attrs.get('coord_start', '2'))
    score_index = int(attrs.get('score_index', '1'))
    id_index = int(attrs.get('id_index', '-1'))
    force_suppress = attrs.get('force_suppress', 'True')
    background_id = int(attrs.get('background_id', '-1'))
    in_format = attrs.get('in_format', 'corner')
    out_format = attrs.get('out_format', 'corner')

    center_point_box = 0 if in_format == 'corner' else 1

    if topk == -1:
        topk = 2**31-1

    if in_format != out_format:
        raise NotImplementedError('box_nms does not currently support in_fomat != out_format')

    if background_id != -1:
        raise NotImplementedError('box_nms does not currently support background_id != -1')

    if id_index != -1 or force_suppress == 'False':
        logging.warning('box_nms: id_idex != -1 or/and force_suppress == False detected. '
                        'However, due to ONNX limitations, boxes of different categories will NOT '
                        'be exempted from suppression. This might lead to different behavior than '
                        'native MXNet')

    create_tensor([coord_start], name+'_cs', kwargs['initializer'])
    create_tensor([coord_start+4], name+'_cs_p4', kwargs['initializer'])
    create_tensor([score_index], name+'_si', kwargs['initializer'])
    create_tensor([score_index+1], name+'_si_p1', kwargs['initializer'])
    create_tensor([topk], name+'_topk', kwargs['initializer'])
    create_tensor([overlap_thresh], name+'_ot', kwargs['initializer'], dtype=np.float32)
    create_tensor([valid_thresh], name+'_vt', kwargs['initializer'], dtype=np.float32)
    create_tensor([-1], name+'_m1', kwargs['initializer'])
    create_tensor([-1], name+'_m1_f', kwargs['initializer'], dtype=dtype)
    create_tensor([0], name+'_0', kwargs['initializer'])
    create_tensor([1], name+'_1', kwargs['initializer'])
    create_tensor([2], name+'_2', kwargs['initializer'])
    create_tensor([3], name+'_3', kwargs['initializer'])
    create_tensor([0, 1, -1], name+'_scores_shape', kwargs['initializer'])
    create_tensor([0, 0, 1, 0], name+'_pad', kwargs['initializer'])
    create_tensor([0, -1], name+'_bat_spat_helper', kwargs['initializer'])
    create_const_scalar_node(name+"_0_s", np.int64(0), kwargs)
    create_const_scalar_node(name+"_1_s", np.int64(1), kwargs)

    nodes = [
        make_node('Shape', [input_nodes[0]], [name+'_shape']),
        make_node('Shape', [name+'_shape'], [name+'_dim']),
        make_node('Sub', [name+'_dim', name+'_2'], [name+'_dim_m2']),
        make_node('Slice', [name+'_shape', name+'_dim_m2', name+'_dim'], [name+'_shape_last2']),
        make_node('Concat', [name+'_m1', name+'_shape_last2'], [name+'_shape_3d'], axis=0),
        make_node('Reshape', [input_nodes[0], name+'_shape_3d'], [name+'_data_3d']),
        make_node('Slice', [name+'_data_3d', name+'_cs', name+'_cs_p4', name+'_m1'],
                  [name+'_boxes']),
        make_node('Slice', [name+'_data_3d', name+'_si', name+'_si_p1', name+'_m1'],
                  [name+'_scores_raw']),
        make_node('Reshape', [name+'_scores_raw', name+'_scores_shape'], [name+'_scores']),
        make_node('Shape', [name+'_scores'], [name+'_scores_shape_actual']),
        make_node('NonMaxSuppression',
                  [name+'_boxes', name+'_scores', name+'_topk', name+'_ot', name+'_vt'],
                  [name+'_nms'], center_point_box=center_point_box),
        make_node('Slice', [name+'_nms', name+'_0', name+'_3', name+'_m1', name+'_2'],
                  [name+'_nms_sliced']),
        make_node('GatherND', [name+'_data_3d', name+'_nms_sliced'], [name+'_candidates']),
        make_node('Pad', [name+'_candidates', name+'_pad', name+'_m1_f'], [name+'_cand_padded']),
        make_node('Shape', [name+'_nms'], [name+'_nms_shape']),
        make_node('Slice', [name+'_nms_shape', name+'_0', name+'_1'], [name+'_cand_cnt']),
        make_node('Squeeze', [name+'_cand_cnt', name+'_0'], [name+'_cc_s']),
        make_node('Range', [name+'_0_s', name+'_cc_s', name+'_1_s'], [name+'_cand_indices']),
        make_node('Slice', [name+'_scores_shape_actual', name+'_0', name+'_3', name+'_m1',
                            name+'_2'], [name+'_shape_bat_spat']),
        make_node('Slice', [name+'_shape_bat_spat', name+'_1', name+'_2'], [name+'_spat_dim']),
        make_node('Expand', [name+'_cand_cnt', name+'_shape_bat_spat'], [name+'_base_indices']),
        make_node('ScatterND', [name+'_base_indices', name+'_nms_sliced', name+'_cand_indices'],
                  [name+'_indices']),
        make_node('TopK', [name+'_indices', name+'_spat_dim'], [name+'_indices_sorted', name+'__'],
                  largest=0, axis=-1, sorted=1),
        make_node('Gather', [name+'_cand_padded', name+'_indices_sorted'], [name+'_gather']),
        make_node('Reshape', [name+'_gather', name+'_shape'], [name+'0'])
    ]

    return nodes


@mx_op.register('_contrib_ROIAlign', OPSET_VERSION)
def convert_contrib_roialign(node, **kwargs):
    """Map MXNet's _contrib_ROIAlign
    """
    from onnx.helper import make_node
    from onnx import TensorProto
    name, input_nodes, attrs = get_inputs(node, kwargs)

    pooled_size = convert_string_to_list(str(attrs.get('pooled_size')))
    spatial_scale = float(attrs.get('spatial_scale'))
    sample_ratio = int(attrs.get('sample_ratio', '0'))
    position_sensitive = attrs.get('position_sensitive', 'False')
    aligned = attrs.get('aligned', 'False')

    if position_sensitive != 'False':
        raise NotImplementedError('_contrib_ROIAlign does not currently support \
                                   position_sensitive!=False')
    if aligned != 'False':
        raise NotImplementedError('_contrib_ROIAlign does not currently support \
                                   aligned!=False')

    create_tensor([0], name+'_0', kwargs['initializer'])
    create_tensor([0], name+'_0_s', kwargs['initializer'], dtype='float32')
    create_tensor([1], name+'_1', kwargs['initializer'])
    create_tensor([5], name+'_5', kwargs['initializer'])
    create_tensor([2, 3], name+'_2_3', kwargs['initializer'])

    nodes = [
        make_node('Slice', [input_nodes[1], name+'_1', name+'_5', name+'_1'], [name+'_rois']),
        make_node('Slice', [input_nodes[1], name+'_0', name+'_1', name+'_1'], [name+'_inds___']),
        make_node('Squeeze', [name+'_inds___', name+'_1'], [name+'_inds__']),
        make_node('Relu', [name+'_inds__'], [name+'_inds_']),
        make_node('Cast', [name+'_inds_'], [name+'_inds'], to=int(TensorProto.INT64)),
        make_node('RoiAlign', [input_nodes[0], name+'_rois', name+'_inds'], [name+'_roi'],
                  mode='avg', output_height=pooled_size[0], output_width=pooled_size[1],
                  sampling_ratio=sample_ratio, spatial_scale=spatial_scale),
        make_node('Unsqueeze', [name+'_inds___', name+'_2_3'], [name+'_unsq']),
        make_node('Less', [name+'_unsq', name+'_0_s'], [name+'_less']),
        make_node('Where', [name+'_less', name+'_0_s', name+'_roi'], [name])
    ]

    return nodes


@mx_op.register("sum", OPSET_VERSION)
def convert_sum(node, **kwargs):
    """Map MXNet's sum operator attributes to onnx's ReduceSum operator
    and return the created node.
    """
    from onnx.helper import make_node

    name, input_nodes, attrs = get_inputs(node, kwargs)

    mx_axis = attrs.get("axis", None)
    axes = convert_string_to_list(str(mx_axis)) if mx_axis is not None else None

    keepdims = get_boolean_attribute_value(attrs, "keepdims")

    if axes:
        create_tensor(axes, name+'_axes', kwargs['initializer'])
        input_nodes.append(name+'_axes')
        node = make_node(
            'ReduceSum',
            inputs=input_nodes,
            outputs=[name],
            keepdims=keepdims,
            name=name
        )
        return [node]
    else:
        create_tensor([1], name+'_1', kwargs['initializer'])
        nodes = [
            onnx.helper.make_node(
                'ReduceSum',
                inputs=input_nodes,
                outputs=[name+'_sum'],
                keepdims=keepdims,
            ),
            make_node('Reshape', [name+'_sum', name+'_1'], [name], name=name),
        ]
    return nodes


@mx_op.register("RNN", OPSET_VERSION)
def convert_RNN(node, **kwargs):
    """Map MXNet's RNN operator attributes to onnx's operators
    and return the created node.
    """
    from onnx.helper import make_node, make_tensor
    from onnx import TensorProto

    name, input_nodes, attrs = get_inputs(node, kwargs)

    mode = str(attrs.get('mode'))
    bidirectional = str(attrs.get('bidirectional', 'False'))
    if bidirectional != 'False' and mode not in ['lstm']:
        raise NotImplementedError('Currently RNN onnx export only supports bidirectional is False')

    num_layers = int(attrs.get('num_layers', '1'))

    use_sequence_length = str(attrs.get('use_sequence_length', 'False'))
    if use_sequence_length != 'False':
        raise NotImplementedError('Currently RNN onnx export only supports use_sequence_length equals to False')

    projection_size = str(attrs.get('projection_size', 'None'))
    if projection_size != 'None':
        raise NotImplementedError('Currently RNN onnx export only supports projection_size equals to None')

    state_outputs = str(attrs.get('state_outputs', 'False'))
    if state_outputs != 'True':
        raise NotImplementedError('Currently RNN onnx export only supports state_outputs equals to True')

    state_size = int(attrs.get('state_size'))

    direction = 1
    if bidirectional != 'False':
        direction = 2

    data = input_nodes[0]
    param = input_nodes[1]
    dtype = get_input_dtypes(node, kwargs)[2]

    create_tensor([0], name+'_0', kwargs['initializer'])
    create_tensor([1], name+'_1', kwargs['initializer'])
    create_tensor([state_size], name+'_state_size', kwargs['initializer'])
    create_tensor([direction], name+'_direction', kwargs['initializer'])

    tensor_1 = make_tensor(name+'_1_f', onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[dtype], [1], [1])

    nodes = [
        make_node('Shape', [data], [name+'_data_shape']),
        make_node('Split', [name+'_data_shape'], [name+'_seq_length', name+'_batch_size', name+'_input_size']),
        make_node('Concat', [name+'_direction', name+'_batch_size', name+'_state_size'], [name+'_concat'], axis=0),
        make_node('ConstantOfShape', [name+'_concat'], [name+'_COS'], value=tensor_1),
        make_node('Mul', [input_nodes[2], name+'_COS'], [name+'initial_h']),

    ]

    if mode == 'lstm':
        nodes += [
            make_node('Mul', [input_nodes[3], name+'_COS'], [name+'initial_c']),
        ]

        if num_layers == 2:
            if bidirectional != 'False':
                raise NotImplementedError('Currently lstm onnx export only supports bidirectional when num_layers = 1')
            create_tensor([8*state_size], name+'_8*state_size', kwargs['initializer'])
            create_tensor([4*state_size*state_size], name+'_4*state_size^2', kwargs['initializer'])
            create_tensor([1, 4*state_size, state_size], name+'_WR_shape', kwargs['initializer'])
            create_tensor([1, 8*state_size], name+'_B_shape', kwargs['initializer'])
            create_tensor([4*4*state_size*state_size], name+'_WR_offset', kwargs['initializer'])

            nodes += [
                # Layer 0
                # get W
                make_node('Slice', [param, name+'_0', name+'_4*state_size^2'], [name+'_W0_1d']),
                make_node('Split', [name+'_W0_1d'], [name+'_W00', name+'_W01', name+'_W02', name+'_W03']),
                make_node('Concat', [name+'_W00', name+'_W03', name+'_W01', name+'_W02'], [name+'_W0_'], axis=0),
                make_node('Reshape', [name+'_W0_', name+'_WR_shape'], [name+'_W0']),
                # get R
                make_node('Add', [name+'_4*state_size^2', name+'_4*state_size^2'], [name+'_R0_offset']),
                make_node('Slice', [param, name+'_4*state_size^2', name+'_R0_offset'], [name+'_R0_1d']),
                make_node('Split', [name+'_R0_1d'], [name+'_R00', name+'_R01', name+'_R02', name+'_R03']),
                make_node('Concat', [name+'_R00', name+'_R03', name+'_R01', name+'_R02'], [name+'_R0_'], axis=0),
                make_node('Reshape', [name+'_R0_', name+'_WR_shape'], [name+'_R0']),
                # get B
                make_node('Add', [name+'_WR_offset', name+'_8*state_size'], [name+'_B0_offset']),
                make_node('Slice', [param, name+'_WR_offset', name+'_B0_offset'], [name+'_B0_1d']),
                make_node('Split', [name+'_B0_1d'], [name+'_B00', name+'_B01', name+'_B02', name+'_B03',
                                                     name+'_B04', name+'_B05', name+'_B06', name+'_B07']),
                make_node('Concat', [name+'_B00', name+'_B03', name+'_B01', name+'_B02',
                                     name+'_B04', name+'_B07', name+'_B05', name+'_B06'], [name+'_B0_'], axis=0),
                make_node('Reshape', [name+'_B0_', name+'_B_shape'], [name+'_B0']),
                # get initial states
                make_node('Split', [name+'initial_h'], [name+'_initial_h0', name+'_initial_h1'], axis=0),
                make_node('Split', [name+'initial_c'], [name+'_initial_c0', name+'_initial_c1'], axis=0),
                # get seq_len
                make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
                make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
                # Layer 0 LSTM
                make_node('LSTM', [data, name+'_W0', name+'_R0', name+'_B0', name+'_seq_len',
                                   name+'_initial_h0', name+'_initial_c0'],
                          [name+'_lstm0_out_', name+'_lstm0_h', name+'_lstm0_c'], hidden_size=state_size),
                make_node('Squeeze', [name+'_lstm0_out_', name+'_1'], [name+'_lstm0_out']),

                # Layer 1
                # get W
                make_node('Add', [name+'_R0_offset', name+'_4*state_size^2'], [name+'_W1_offset']),
                make_node('Slice', [param, name+'_R0_offset', name+'_W1_offset'], [name+'_W1_1d']),
                make_node('Split', [name+'_W1_1d'], [name+'_W10', name+'_W11', name+'_W12', name+'_W13']),
                make_node('Concat', [name+'_W10', name+'_W13', name+'_W11', name+'_W12'], [name+'_W1_'], axis=0),
                make_node('Reshape', [name+'_W1_', name+'_WR_shape'], [name+'_W1']),
                # get R
                make_node('Slice', [param, name+'_W1_offset', name+'_WR_offset'], [name+'_R1_1d']),
                make_node('Split', [name+'_R1_1d'], [name+'_R10', name+'_R11', name+'_R12', name+'_R13']),
                make_node('Concat', [name+'_R10', name+'_R13', name+'_R11', name+'_R12'], [name+'_R1_'], axis=0),
                make_node('Reshape', [name+'_R1_', name+'_WR_shape'], [name+'_R1']),
                # get B
                make_node('Add', [name+'_B0_offset', name+'_8*state_size'], [name+'_B1_offset']),
                make_node('Slice', [param, name+'_B0_offset', name+'_B1_offset'], [name+'_B1_1d']),
                make_node('Split', [name+'_B1_1d'], [name+'_B10', name+'_B11', name+'_B12', name+'_B13',
                                                     name+'_B14', name+'_B15', name+'_B16', name+'_B17']),
                make_node('Concat', [name+'_B10', name+'_B13', name+'_B11', name+'_B12',
                                     name+'_B14', name+'_B17', name+'_B15', name+'_B16'], [name+'_B1_'], axis=0),
                make_node('Reshape', [name+'_B1_', name+'_B_shape'], [name+'_B1']),
                # Layer 1 LSTM
                make_node('LSTM', [name+'_lstm0_out', name+'_W1', name+'_R1', name+'_B1', name+'_seq_len',
                                   name+'_initial_h1', name+'_initial_c1'],
                          [name+'_lstm1_out_', name+'_lstm1_h', name+'_lstm1_c'], hidden_size=state_size),
                make_node('Squeeze', [name+'_lstm1_out_', name+'_1'], [name]),
                make_node('Concat', [name+'_lstm0_h', name+'_lstm1_h'], [name+'1'], axis=0),
                make_node('Concat', [name+'_lstm0_c', name+'_lstm1_c'], [name+'2'], axis=0),
            ]
        elif num_layers == 1:
            if bidirectional == 'False':
                create_tensor([4*state_size], name+'_4*state_size', kwargs['initializer'])
                create_tensor([8*state_size], name+'_8*state_size', kwargs['initializer'])
                create_tensor([4*state_size*state_size], name+'_4*state_size^2', kwargs['initializer'])
                create_tensor([1, 4*state_size, state_size], name+'_R_shape', kwargs['initializer'])
                create_tensor([1, 8*state_size], name+'_B_shape', kwargs['initializer'])

                nodes += [
                    # get W
                    make_node('Mul', [name+'_4*state_size', name+'_input_size'], [name+'_mul0']),
                    make_node('Slice', [param, name+'_0', name+'_mul0'], [name+'_W_1d']),
                    make_node('Split', [name+'_W_1d'], [name+'_W0', name+'_W1', name+'_W2', name+'_W3']),
                    make_node('Concat', [name+'_W0', name+'_W3', name+'_W1', name+'_W2'], [name+'_W_'], axis=0),
                    make_node('Concat', [name+'_1', name+'_4*state_size', name+'_input_size'],
                              [name+'_W_shape'], axis=0),
                    make_node('Reshape', [name+'_W_', name+'_W_shape'], [name+'_W']),
                    # get R
                    make_node('Add', [name+'_mul0', name+'_4*state_size^2'], [name+'_add0']),
                    make_node('Slice', [param, name+'_mul0', name+'_add0'], [name+'_R_1d']),
                    make_node('Split', [name+'_R_1d'], [name+'_R0', name+'_R1', name+'_R2', name+'_R3']),
                    make_node('Concat', [name+'_R0', name+'_R3', name+'_R1', name+'_R2'], [name+'_R_'], axis=0),
                    make_node('Reshape', [name+'_R_', name+'_R_shape'], [name+'_R']),
                    # get B
                    make_node('Add', [name+'_add0', name+'_8*state_size'], [name+'_add1']),
                    make_node('Slice', [param, name+'_add0', name+'_add1'], [name+'_B_1d']),
                    make_node('Split', [name+'_B_1d'], [name+'_B0', name+'_B1', name+'_B2', name+'_B3',
                                                        name+'_B4', name+'_B5', name+'_B6', name+'_B7']),
                    make_node('Concat', [name+'_B0', name+'_B3', name+'_B1', name+'_B2',
                                         name+'_B4', name+'_B7', name+'_B5', name+'_B6'], [name+'_B_'], axis=0),
                    make_node('Reshape', [name+'_B_', name+'_B_shape'], [name+'_B']),
                    # get seq_len
                    make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
                    make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
                    # compute LSTM
                    make_node('LSTM', [data, name+'_W', name+'_R', name+'_B',
                                       name+'_seq_len', name+'initial_h', name+'initial_c'],
                              [name+'0_', name+'1', name+'2'], hidden_size=state_size),
                    make_node('Squeeze', [name+'0_', name+'_1'], [name]),
                ]
            else:
                create_tensor([-1], name+'_-1', kwargs['initializer'])
                create_tensor([4*state_size], name+'_4*state_size', kwargs['initializer'])
                create_tensor([8*state_size], name+'_8*state_size', kwargs['initializer'])
                create_tensor([4*state_size*state_size], name+'_4*state_size^2', kwargs['initializer'])
                create_tensor([1, 4*state_size, state_size], name+'_R_shape', kwargs['initializer'])
                create_tensor([1, 8*state_size], name+'_B_shape', kwargs['initializer'])

                nodes += [
                    # get W_fwd
                    make_node('Mul', [name+'_4*state_size', name+'_input_size'], [name+'_mul0']),
                    make_node('Slice', [param, name+'_0', name+'_mul0'], [name+'_W_1d']),
                    make_node('Split', [name+'_W_1d'], [name+'_W0', name+'_W1', name+'_W2', name+'_W3']),
                    make_node('Concat', [name+'_W0', name+'_W3', name+'_W1', name+'_W2'],
                              [name+'_W_'], axis=0),
                    make_node('Concat', [name+'_1', name+'_4*state_size', name+'_input_size'],
                              [name+'_W_shape'], axis=0),
                    make_node('Reshape', [name+'_W_', name+'_W_shape'], [name+'_W_fwd']),
                    # get R_fwd
                    make_node('Add', [name+'_mul0', name+'_4*state_size^2'], [name+'_add0']),
                    make_node('Slice', [param, name+'_mul0', name+'_add0'], [name+'_R_1d']),
                    make_node('Split', [name+'_R_1d'], [name+'_R0', name+'_R1', name+'_R2', name+'_R3']),
                    make_node('Concat', [name+'_R0', name+'_R3', name+'_R1', name+'_R2'], [name+'_R_'], axis=0),
                    make_node('Reshape', [name+'_R_', name+'_R_shape'], [name+'_R_fwd']),
                    # get W_bwd
                    make_node('Add', [name+'_add0', name+'_mul0'], [name+'_add1']),
                    make_node('Slice', [param, name+'_add0', name+'_add1'], [name+'_W_1d_bwd']),
                    make_node('Split', [name+'_W_1d_bwd'],
                              [name+'_W0_bwd', name+'_W1_bwd', name+'_W2_bwd', name+'_W3_bwd']),
                    make_node('Concat', [name+'_W0_bwd', name+'_W3_bwd', name+'_W1_bwd', name+'_W2_bwd'],
                              [name+'_W_bwd_'], axis=0),
                    make_node('Reshape', [name+'_W_bwd_', name+'_W_shape'], [name+'_W_bwd']),
                    # get R_bwd
                    make_node('Add', [name+'_add1', name+'_4*state_size^2'], [name+'_add2']),
                    make_node('Slice', [param, name+'_add1', name+'_add2'], [name+'_R_1d_bwd']),
                    make_node('Split', [name+'_R_1d_bwd'],
                              [name+'_R0_bwd', name+'_R1_bwd', name+'_R2_bwd', name+'_R3_bwd']),
                    make_node('Concat', [name+'_R0_bwd', name+'_R3_bwd', name+'_R1_bwd', name+'_R2_bwd'],
                              [name+'_R_bwd_'], axis=0),
                    make_node('Reshape', [name+'_R_bwd_', name+'_R_shape'], [name+'_R_bwd']),
                    # get B_fwd
                    make_node('Add', [name+'_add2', name+'_8*state_size'], [name+'_add3']),
                    make_node('Slice', [param, name+'_add2', name+'_add3'], [name+'_B_1d']),
                    make_node('Split', [name+'_B_1d'], [name+'_B0', name+'_B1', name+'_B2', name+'_B3',
                                                        name+'_B4', name+'_B5', name+'_B6', name+'_B7']),
                    make_node('Concat', [name+'_B0', name+'_B3', name+'_B1', name+'_B2',
                                         name+'_B4', name+'_B7', name+'_B5', name+'_B6'], [name+'_B_'], axis=0),
                    make_node('Reshape', [name+'_B_', name+'_B_shape'], [name+'_B_fwd']),
                    # get B_bwd
                    make_node('Add', [name+'_add3', name+'_8*state_size'], [name+'_add4']),
                    make_node('Slice', [param, name+'_add3', name+'_add4'], [name+'_B_1d_bwd']),
                    make_node('Split', [name+'_B_1d_bwd'],
                              [name+'_B0_bwd', name+'_B1_bwd', name+'_B2_bwd', name+'_B3_bwd',
                               name+'_B4_bwd', name+'_B5_bwd', name+'_B6_bwd', name+'_B7_bwd']),
                    make_node('Concat', [name+'_B0_bwd', name+'_B3_bwd', name+'_B1_bwd', name+'_B2_bwd',
                                         name+'_B4_bwd', name+'_B7_bwd', name+'_B5_bwd', name+'_B6_bwd'],
                              [name+'_B_bwd_'], axis=0),
                    make_node('Reshape', [name+'_B_bwd_', name+'_B_shape'], [name+'_B_bwd']),
                    # get seq_len
                    make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
                    make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
                    # compute LSTM
                    make_node('Concat', [name+'_W_fwd', name+'_W_bwd'], [name+'_W'], axis=0),
                    make_node('Concat', [name+'_R_fwd', name+'_R_bwd'], [name+'_R'], axis=0),
                    make_node('Concat', [name+'_B_fwd', name+'_B_bwd'], [name+'_B'], axis=0),
                    make_node('LSTM', [data, name+'_W', name+'_R', name+'_B',
                                       name+'_seq_len', name+'initial_h', name+'initial_c'],
                              [name+'0_', name+'1', name+'2'], hidden_size=state_size, direction='bidirectional'),
                    make_node('Transpose', [name+'0_'], [name+'0_t'], perm=[0, 2, 1, 3]),
                    make_node('Concat', [name+'_seq_length', name+'_batch_size', name+'_-1'],
                              [name+'_shape_out'], axis=0),
                    make_node('Reshape', [name+'0_t', name+'_shape_out'], [name]),
                ]
        else:
            raise NotImplementedError('Currently RNN onnx export only supports num_layers equals to 1 or 2')

    elif mode == 'gru':
        if num_layers == 2:
            create_tensor([6*state_size], name+'_6*state_size', kwargs['initializer'])
            create_tensor([3*state_size*state_size], name+'_3*state_size^2', kwargs['initializer'])
            create_tensor([1, 3*state_size, state_size], name+'_WR_shape', kwargs['initializer'])
            create_tensor([1, 6*state_size], name+'_B_shape', kwargs['initializer'])
            create_tensor([4*3*state_size*state_size], name+'_WR_offset', kwargs['initializer'])

            nodes += [
                # Layer 0
                # get W
                make_node('Slice', [param, name+'_0', name+'_3*state_size^2'], [name+'_W0_1d']),
                make_node('Split', [name+'_W0_1d'], [name+'_W00', name+'_W01', name+'_W02']),
                make_node('Concat', [name+'_W01', name+'_W00', name+'_W02'], [name+'_W0_'], axis=0),
                make_node('Reshape', [name+'_W0_', name+'_WR_shape'], [name+'_W0']),
                # get R
                make_node('Add', [name+'_3*state_size^2', name+'_3*state_size^2'], [name+'_R0_offset']),
                make_node('Slice', [param, name+'_3*state_size^2', name+'_R0_offset'], [name+'_R0_1d']),
                make_node('Split', [name+'_R0_1d'], [name+'_R00', name+'_R01', name+'_R02']),
                make_node('Concat', [name+'_R01', name+'_R00', name+'_R02'], [name+'_R0_'], axis=0),
                make_node('Reshape', [name+'_R0_', name+'_WR_shape'], [name+'_R0']),
                # get B
                make_node('Add', [name+'_WR_offset', name+'_6*state_size'], [name+'_B0_offset']),
                make_node('Slice', [param, name+'_WR_offset', name+'_B0_offset'], [name+'_B0_1d']),
                make_node('Split', [name+'_B0_1d'], [name+'_B00', name+'_B01', name+'_B02',
                                                     name+'_B03', name+'_B04', name+'_B05']),
                make_node('Concat', [name+'_B01', name+'_B00', name+'_B02',
                                     name+'_B04', name+'_B03', name+'_B05'], [name+'_B0_'], axis=0),
                make_node('Reshape', [name+'_B0_', name+'_B_shape'], [name+'_B0']),
                # get initial states
                make_node('Split', [name+'initial_h'], [name+'_initial_h0', name+'_initial_h1'], axis=0),
                # get seq_len
                make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
                make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
                # Layer 0 GRU
                make_node('GRU', [data, name+'_W0', name+'_R0', name+'_B0', name+'_seq_len',
                                  name+'_initial_h0'],
                          [name+'_gru0_out_', name+'_gru0_h'], hidden_size=state_size, linear_before_reset=1),
                make_node('Squeeze', [name+'_gru0_out_', name+'_1'], [name+'_gru0_out']),

                # Layer 1
                # get W
                make_node('Add', [name+'_R0_offset', name+'_3*state_size^2'], [name+'_W1_offset']),
                make_node('Slice', [param, name+'_R0_offset', name+'_W1_offset'], [name+'_W1_1d']),
                make_node('Split', [name+'_W1_1d'], [name+'_W10', name+'_W11', name+'_W12']),
                make_node('Concat', [name+'_W11', name+'_W10', name+'_W12'], [name+'_W1_'], axis=0),
                make_node('Reshape', [name+'_W1_', name+'_WR_shape'], [name+'_W1']),
                # get R
                make_node('Slice', [param, name+'_W1_offset', name+'_WR_offset'], [name+'_R1_1d']),
                make_node('Split', [name+'_R1_1d'], [name+'_R10', name+'_R11', name+'_R12']),
                make_node('Concat', [name+'_R11', name+'_R10', name+'_R12'], [name+'_R1_'], axis=0),
                make_node('Reshape', [name+'_R1_', name+'_WR_shape'], [name+'_R1']),
                # get B
                make_node('Add', [name+'_B0_offset', name+'_6*state_size'], [name+'_B1_offset']),
                make_node('Slice', [param, name+'_B0_offset', name+'_B1_offset'], [name+'_B1_1d']),
                make_node('Split', [name+'_B1_1d'], [name+'_B10', name+'_B11', name+'_B12',
                                                     name+'_B13', name+'_B14', name+'_B15']),
                make_node('Concat', [name+'_B11', name+'_B10', name+'_B12',
                                     name+'_B14', name+'_B13', name+'_B15'], [name+'_B1_'], axis=0),
                make_node('Reshape', [name+'_B1_', name+'_B_shape'], [name+'_B1']),
                # Layer 1 GRU
                make_node('GRU', [name+'_gru0_out', name+'_W1', name+'_R1', name+'_B1', name+'_seq_len',
                                  name+'_initial_h1'],
                          [name+'_gru1_out_', name+'_gru1_h'], hidden_size=state_size, linear_before_reset=1),
                make_node('Squeeze', [name+'_gru1_out_', name+'_1'], [name]),
                make_node('Concat', [name+'_gru0_h', name+'_gru1_h'], [name+'1'], axis=0)
            ]

        elif num_layers == 1:
            create_tensor([3*state_size], name+'_3*state_size', kwargs['initializer'])
            create_tensor([6*state_size], name+'_6*state_size', kwargs['initializer'])
            create_tensor([3*state_size*state_size], name+'_3*state_size^2', kwargs['initializer'])
            create_tensor([1, 3*state_size, state_size], name+'_R_shape', kwargs['initializer'])
            create_tensor([1, 6*state_size], name+'_B_shape', kwargs['initializer'])

            nodes += [
                # get W
                make_node('Mul', [name+'_3*state_size', name+'_input_size'], [name+'_mul0']),
                make_node('Slice', [param, name+'_0', name+'_mul0'], [name+'_W_1d']),
                make_node('Split', [name+'_W_1d'], [name+'_W0', name+'_W1', name+'_W2']),
                make_node('Concat', [name+'_W1', name+'_W0', name+'_W2'], [name+'_W_'], axis=0),
                make_node('Concat', [name+'_1', name+'_3*state_size', name+'_input_size'], [name+'_W_shape'], axis=0),
                make_node('Reshape', [name+'_W_', name+'_W_shape'], [name+'_W']),
                # get R
                make_node('Add', [name+'_mul0', name+'_3*state_size^2'], [name+'_add0']),
                make_node('Slice', [param, name+'_mul0', name+'_add0'], [name+'_R_1d']),
                make_node('Split', [name+'_R_1d'], [name+'_R0', name+'_R1', name+'_R2']),
                make_node('Concat', [name+'_R1', name+'_R0', name+'_R2'], [name+'_R_'], axis=0),
                make_node('Reshape', [name+'_R_', name+'_R_shape'], [name+'_R']),
                # get B
                make_node('Add', [name+'_add0', name+'_6*state_size'], [name+'_add1']),
                make_node('Slice', [param, name+'_add0', name+'_add1'], [name+'_B_1d']),
                make_node('Split', [name+'_B_1d'], [name+'_B0', name+'_B1', name+'_B2',
                                                    name+'_B3', name+'_B4', name+'_B5']),
                make_node('Concat', [name+'_B1', name+'_B0', name+'_B2',
                                     name+'_B4', name+'_B3', name+'_B5'], [name+'_B_'], axis=0),
                make_node('Reshape', [name+'_B_', name+'_B_shape'], [name+'_B']),
                # get seq_len
                make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
                make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
                # compute GRU
                make_node('GRU', [data, name+'_W', name+'_R', name+'_B', name+'_seq_len', name+'initial_h'],
                          [name+'0_', name+'1'], hidden_size=state_size, linear_before_reset=1),
                make_node('Squeeze', [name+'0_', name+'_1'], [name]),
            ]
        else:
            raise NotImplementedError('Currently RNN onnx export only supports num_layers equals to 1 or 2')

    elif mode in ['rnn_tanh', 'rnn_relu']:
        activations = ['Tanh']
        if mode == 'rnn_relu':
            activations = ['Relu']
        if num_layers == 2:
            create_tensor([2*state_size], name+'_2*state_size', kwargs['initializer'])
            create_tensor([state_size*state_size], name+'_state_size^2', kwargs['initializer'])
            create_tensor([1, state_size, state_size], name+'_WR_shape', kwargs['initializer'])
            create_tensor([1, 2*state_size], name+'_B_shape', kwargs['initializer'])
            create_tensor([4*state_size*state_size], name+'_WR_offset', kwargs['initializer'])

            nodes += [
                # Layer 0
                # get W
                make_node('Slice', [param, name+'_0', name+'_state_size^2'], [name+'_W0_1d']),
                make_node('Reshape', [name+'_W0_1d', name+'_WR_shape'], [name+'_W0']),
                # get R
                make_node('Add', [name+'_state_size^2', name+'_state_size^2'], [name+'_R0_offset']),
                make_node('Slice', [param, name+'_state_size^2', name+'_R0_offset'], [name+'_R0_1d']),
                make_node('Reshape', [name+'_R0_1d', name+'_WR_shape'], [name+'_R0']),
                # get B
                make_node('Add', [name+'_WR_offset', name+'_2*state_size'], [name+'_B0_offset']),
                make_node('Slice', [param, name+'_WR_offset', name+'_B0_offset'], [name+'_B0_1d']),
                make_node('Reshape', [name+'_B0_1d', name+'_B_shape'], [name+'_B0']),
                # get initial states
                make_node('Split', [name+'initial_h'], [name+'_initial_h0', name+'_initial_h1'], axis=0),
                # get seq_len
                make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
                make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
                # Layer 0 RNN
                make_node('RNN', [data, name+'_W0', name+'_R0', name+'_B0', name+'_seq_len',
                                  name+'_initial_h0'], [name+'_rnn0_out_', name+'_rnn0_h'],
                          hidden_size=state_size, activations=activations),
                make_node('Squeeze', [name+'_rnn0_out_', name+'_1'], [name+'_rnn0_out']),

                # Layer 1
                # get W
                make_node('Add', [name+'_R0_offset', name+'_state_size^2'], [name+'_W1_offset']),
                make_node('Slice', [param, name+'_R0_offset', name+'_W1_offset'], [name+'_W1_1d']),
                make_node('Reshape', [name+'_W1_1d', name+'_WR_shape'], [name+'_W1']),
                # get R
                make_node('Slice', [param, name+'_W1_offset', name+'_WR_offset'], [name+'_R1_1d']),
                make_node('Reshape', [name+'_R1_1d', name+'_WR_shape'], [name+'_R1']),
                # get B
                make_node('Add', [name+'_B0_offset', name+'_2*state_size'], [name+'_B1_offset']),
                make_node('Slice', [param, name+'_B0_offset', name+'_B1_offset'], [name+'_B1_1d']),
                make_node('Reshape', [name+'_B1_1d', name+'_B_shape'], [name+'_B1']),
                # Layer 1 RNN
                make_node('RNN', [name+'_rnn0_out', name+'_W1', name+'_R1', name+'_B1', name+'_seq_len',
                                  name+'_initial_h1'], [name+'_rnn1_out_', name+'_rnn1_h'],
                          hidden_size=state_size, activations=activations),
                make_node('Squeeze', [name+'_rnn1_out_', name+'_1'], [name]),
                make_node('Concat', [name+'_rnn0_h', name+'_rnn1_h'], [name+'1'], axis=0)
            ]

        elif num_layers == 1:
            create_tensor([2*state_size], name+'_2*state_size', kwargs['initializer'])
            create_tensor([state_size*state_size], name+'_state_size^2', kwargs['initializer'])
            create_tensor([1, state_size, state_size], name+'_R_shape', kwargs['initializer'])
            create_tensor([1, 2*state_size], name+'_B_shape', kwargs['initializer'])

            nodes += [
                # get W
                make_node('Mul', [name+'_state_size', name+'_input_size'], [name+'_mul0']),
                make_node('Slice', [param, name+'_0', name+'_mul0'], [name+'_W_1d']),
                make_node('Concat', [name+'_1', name+'_state_size', name+'_input_size'], [name+'_W_shape'], axis=0),
                make_node('Reshape', [name+'_W_1d', name+'_W_shape'], [name+'_W']),
                # get R
                make_node('Add', [name+'_mul0', name+'_state_size^2'], [name+'_add0']),
                make_node('Slice', [param, name+'_mul0', name+'_add0'], [name+'_R_1d']),
                make_node('Reshape', [name+'_R_1d', name+'_R_shape'], [name+'_R']),
                # get B
                make_node('Add', [name+'_add0', name+'_2*state_size'], [name+'_add1']),
                make_node('Slice', [param, name+'_add0', name+'_add1'], [name+'_B_1d']),
                make_node('Reshape', [name+'_B_1d', name+'_B_shape'], [name+'_B']),
                # get seq_len
                make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
                make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
                # compute RNN
                make_node('RNN', [data, name+'_W', name+'_R', name+'_B', name+'_seq_len', name+'initial_h'],
                          [name+'0_', name+'1'], hidden_size=state_size, activations=activations),
                make_node('Squeeze', [name+'0_', name+'_1'], [name]),
            ]
        else:
            raise NotImplementedError('Currently RNN onnx export only supports num_layers equals to 1 or 2')
    else:
        raise NotImplementedError(f"Currently RNN onnx export does not support {mode} mode")
    return nodes


@mx_op.register('SliceChannel', OPSET_VERSION)
def convert_slice_channel(node, **kwargs):
    """Map MXNet's SliceChannel operator attributes to onnx's Squeeze or Split
    operator based on squeeze_axis attribute
    and return the created node.
    """
    from onnx.helper import make_node
    name, input_nodes, attrs = get_inputs(node, kwargs)

    num_outputs = int(attrs.get('num_outputs'))
    axis = int(attrs.get('axis', 1))
    squeeze_axis = attrs.get('squeeze_axis', 'False')

    create_tensor([axis], name+'_axis', kwargs['initializer'])

    nodes = []
    if squeeze_axis in ['True', '1']:
        nodes += [
            make_node('Split', [input_nodes[0]], [name+str(i)+'_' for i in range(num_outputs)],
                      axis=axis)
        ]
        for i in range(num_outputs):
            nodes += [
                make_node('Squeeze', [name+str(i)+'_', name+'_axis'], [name+str(i)])
            ]
    else:
        nodes += [
            make_node('Split', [input_nodes[0]], [name+str(i) for i in range(num_outputs)],
                      axis=axis)
        ]

    return nodes


@mx_op.register("max", OPSET_VERSION)
def convert_max(node, **kwargs):
    """Map MXNet's max operator attributes to onnx's ReduceMax operator
    and return the created node.
    """
    from onnx.helper import make_node
    name, input_nodes, attrs = get_inputs(node, kwargs)

    mx_axis = str(attrs.get("axis", 'None'))
    axes = convert_string_to_list(mx_axis) if mx_axis != 'None' else None

    keepdims = get_boolean_attribute_value(attrs, "keepdims")

    if axes is not None:
        if keepdims:
            node = make_node('ReduceMax', input_nodes, [name], axes=axes, keepdims=keepdims)
            return [node]
        else:
            create_tensor([1], name+'_1', kwargs['initializer'])
            create_tensor([0], name+'_0', kwargs['initializer'])
            create_tensor([len(axes)], name+'_axes_dim', kwargs['initializer'])
            nodes = [
                make_node('ReduceMax', input_nodes, [name+'_rmax'], axes=axes, keepdims=keepdims),
                make_node('Shape', [name+'_rmax'], [name+'_rmax_shape']),
                make_node('Shape', [name+'_rmax_shape'], [name+'_rmax_dim']),
                make_node('Shape', [input_nodes[0]], [name+'_in_shape']),
                make_node('Shape', [name+'_in_shape'], [name+'_in_dim']),
                make_node('Equal', [name+'_axes_dim', name+'_in_dim'], [name+'_equal']),
                make_node('Where', [name+'_equal', name+'_1', name+'_rmax_dim'], [name+'_where0']),
                make_node('Tile', [name+'_0', name+'_where0'], [name+'_tile']),
                make_node('Unsqueeze', [name+'_0', name+'_0'], [name+'_unsqueeze']),
                make_node('Where', [name+'_equal', name+'_1', name+'_0'], [name+'_where1']),
                make_node('ScatterND', [name+'_tile', name+'_unsqueeze', name+'_where1'], [name+'_SND']),
                make_node('Reshape', [name+'_rmax', name+'_SND'], [name]),
            ]
            return nodes
    else:
        if keepdims:
            node = make_node('ReduceMax', input_nodes, [name], keepdims=keepdims)
            return [node]
        else:
            create_tensor([1], name+'_1', kwargs['initializer'])
            nodes = [
                make_node('ReduceMax', input_nodes, [name+'_rmax'], keepdims=keepdims),
                make_node('Reshape', [name+'_rmax', name+'_1'], [name])
            ]
            return nodes


@mx_op.register("min", OPSET_VERSION)
def convert_min(node, **kwargs):
    """Map MXNet's min operator attributes to onnx's ReduceMin operator
    and return the created node.
    """
    from onnx.helper import make_node
    name, input_nodes, attrs = get_inputs(node, kwargs)

    mx_axis = str(attrs.get("axis", 'None'))
    axes = convert_string_to_list(mx_axis) if mx_axis != 'None' else None

    keepdims = get_boolean_attribute_value(attrs, "keepdims")

    if axes is not None:
        if keepdims:
            node = make_node('ReduceMin', input_nodes, [name], axes=axes, keepdims=keepdims)
            return [node]
        else:
            create_tensor([1], name+'_1', kwargs['initializer'])
            create_tensor([0], name+'_0', kwargs['initializer'])
            create_tensor([len(axes)], name+'_axes_dim', kwargs['initializer'])
            nodes = [
                make_node('ReduceMin', input_nodes, [name+'_rmin'], axes=axes, keepdims=keepdims),
                make_node('Shape', [name+'_rmin'], [name+'_rmin_shape']),
                make_node('Shape', [name+'_rmin_shape'], [name+'_rmin_dim']),
                make_node('Shape', [input_nodes[0]], [name+'_in_shape']),
                make_node('Shape', [name+'_in_shape'], [name+'_in_dim']),
                make_node('Equal', [name+'_axes_dim', name+'_in_dim'], [name+'_equal']),
                make_node('Where', [name+'_equal', name+'_1', name+'_rmin_dim'], [name+'_where0']),
                make_node('Tile', [name+'_0', name+'_where0'], [name+'_tile']),
                make_node('Unsqueeze', [name+'_0', name+'_0'], [name+'_unsqueeze']),
                make_node('Where', [name+'_equal', name+'_1', name+'_0'], [name+'_where1']),
                make_node('ScatterND', [name+'_tile', name+'_unsqueeze', name+'_where1'], [name+'_SND']),
                make_node('Reshape', [name+'_rmin', name+'_SND'], [name]),
            ]
            return nodes
    else:
        if keepdims:
            node = make_node('ReduceMin', input_nodes, [name], keepdims=keepdims)
            return [node]

        else:
            create_tensor([1], name+'_1', kwargs['initializer'])
            nodes = [
                make_node('ReduceMin', input_nodes, [name+'_rmin'], keepdims=keepdims),
                make_node('Reshape', [name+'_rmin', name+'_1'], [name])
            ]
            return nodes


@mx_op.register("mean", OPSET_VERSION)
def convert_mean(node, **kwargs):
    """Map MXNet's mean operator attributes to onnx's ReduceMean operator
    and return the created node.
    """
    from onnx.helper import make_node
    name, input_nodes, attrs = get_inputs(node, kwargs)

    mx_axis = str(attrs.get("axis", 'None'))
    axes = convert_string_to_list(mx_axis) if mx_axis != 'None' else None

    keepdims = get_boolean_attribute_value(attrs, "keepdims")

    if axes is not None:
        if keepdims:
            node = make_node('ReduceMean', input_nodes, [name], axes=axes, keepdims=keepdims)
            return [node]
        else:
            create_tensor([1], name+'_1', kwargs['initializer'])
            create_tensor([0], name+'_0', kwargs['initializer'])
            create_tensor([len(axes)], name+'_axes_dim', kwargs['initializer'])
            nodes = [
                make_node('ReduceMean', input_nodes, [name+'_reduce'], axes=axes, keepdims=keepdims),
                make_node('Shape', [name+'_reduce'], [name+'_reduce_shape']),
                make_node('Shape', [name+'_reduce_shape'], [name+'_reduce_dim']),
                make_node('Shape', [input_nodes[0]], [name+'_in_shape']),
                make_node('Shape', [name+'_in_shape'], [name+'_in_dim']),
                make_node('Equal', [name+'_axes_dim', name+'_in_dim'], [name+'_equal']),
                make_node('Where', [name+'_equal', name+'_1', name+'_reduce_dim'], [name+'_where0']),
                make_node('Tile', [name+'_0', name+'_where0'], [name+'_tile']),
                make_node('Unsqueeze', [name+'_0', name+'_0'], [name+'_unsqueeze']),
                make_node('Where', [name+'_equal', name+'_1', name+'_0'], [name+'_where1']),
                make_node('ScatterND', [name+'_tile', name+'_unsqueeze', name+'_where1'], [name+'_SND']),
                make_node('Reshape', [name+'_reduce', name+'_SND'], [name]),
            ]
            return nodes
    else:
        if keepdims:
            node = make_node('ReduceMean', input_nodes, [name], keepdims=keepdims)
            return [node]

        else:
            create_tensor([1], name+'_1', kwargs['initializer'])
            nodes = [
                make_node('ReduceMean', input_nodes, [name+'_reduce'], keepdims=keepdims),
                make_node('Reshape', [name+'_reduce', name+'_1'], [name])
            ]
            return nodes


@mx_op.register("prod", OPSET_VERSION)
def convert_prod(node, **kwargs):
    """Map MXNet's prod operator attributes to onnx's ReduceProd operator
    and return the created node.
    """
    from onnx.helper import make_node
    name, input_nodes, attrs = get_inputs(node, kwargs)

    mx_axis = str(attrs.get("axis", 'None'))
    axes = convert_string_to_list(mx_axis) if mx_axis != 'None' else None

    keepdims = get_boolean_attribute_value(attrs, "keepdims")

    if axes is not None:
        if keepdims:
            node = make_node('ReduceProd', input_nodes, [name], axes=axes, keepdims=keepdims)
            return [node]
        else:
            create_tensor([1], name+'_1', kwargs['initializer'])
            create_tensor([0], name+'_0', kwargs['initializer'])
            create_tensor([len(axes)], name+'_axes_dim', kwargs['initializer'])
            nodes = [
                make_node('ReduceProd', input_nodes, [name+'_reduce'], axes=axes, keepdims=keepdims),
                make_node('Shape', [name+'_reduce'], [name+'_reduce_shape']),
                make_node('Shape', [name+'_reduce_shape'], [name+'_reduce_dim']),
                make_node('Shape', [input_nodes[0]], [name+'_in_shape']),
                make_node('Shape', [name+'_in_shape'], [name+'_in_dim']),
                make_node('Equal', [name+'_axes_dim', name+'_in_dim'], [name+'_equal']),
                make_node('Where', [name+'_equal', name+'_1', name+'_reduce_dim'], [name+'_where0']),
                make_node('Tile', [name+'_0', name+'_where0'], [name+'_tile']),
                make_node('Unsqueeze', [name+'_0', name+'_0'], [name+'_unsqueeze']),
                make_node('Where', [name+'_equal', name+'_1', name+'_0'], [name+'_where1']),
                make_node('ScatterND', [name+'_tile', name+'_unsqueeze', name+'_where1'], [name+'_SND']),
                make_node('Reshape', [name+'_reduce', name+'_SND'], [name]),
            ]
            return nodes
    else:
        if keepdims:
            node = make_node('ReduceProd', input_nodes, [name], keepdims=keepdims)
            return [node]

        else:
            create_tensor([1], name+'_1', kwargs['initializer'])
            nodes = [
                make_node('ReduceProd', input_nodes, [name+'_reduce'], keepdims=keepdims),
                make_node('Reshape', [name+'_reduce', name+'_1'], [name])
            ]
            return nodes


@mx_op.register("squeeze", OPSET_VERSION)
def convert_squeeze(node, **kwargs):
    """Map MXNet's squeeze operator attributes to onnx's squeeze operator
    and return the created node.
    """
    name, input_nodes, attrs = get_inputs(node, kwargs)

    mx_axis = str(attrs.get("axis", 'None'))
    axes = convert_string_to_list(mx_axis) if mx_axis != 'None' else None

    if not axes:
        node = onnx.helper.make_node(
            "Squeeze",
            input_nodes,
            [name],
            name=name
        )
    else:
        create_tensor(axes, name+'_axes', kwargs['initializer'])
        node = onnx.helper.make_node(
            "Squeeze",
            [input_nodes[0], name+'_axes'],
            [name],
            name=name,
        )
    return [node]


@mx_op.register("SoftmaxOutput", OPSET_VERSION)
def convert_softmax_output(node, **kwargs):
    """Map MXNet's SoftmaxOutput operator attributes to onnx's Softmax operator
    and return the created node.
    """
    from onnx.helper import make_node
    name, input_nodes, _ = get_inputs(node, kwargs)

    nodes = [
        make_node('Shape', [input_nodes[0]], [name+'_shape']),
        make_node('Flatten', [input_nodes[0]], [name+'_flat'], axis=1),
        make_node('Softmax', [name+'_flat'], [name+'_sm'], axis=1),
        make_node('Reshape', [name+'_sm', name+'_shape'], [name])
    ]

    return nodes


@mx_op.register("norm", OPSET_VERSION)
def convert_norm(node, **kwargs):
    """Map MXNet's norm operator attributes to onnx's ReduceL1 and ReduceL2 operators
    and return the created node.
    """
    from onnx.helper import make_node
    name, input_nodes, attrs = get_inputs(node, kwargs)

    mx_axis = attrs.get("axis", None)
    axes = convert_string_to_list(str(mx_axis)) if mx_axis else None

    keepdims = get_boolean_attribute_value(attrs, "keepdims")
    ord = int(attrs.get("ord", 2))

    onnx_op_name = "ReduceL1" if ord == 1 else "ReduceL2"

    if axes:
        if keepdims:
            reduce_node = make_node(onnx_op_name, input_nodes, [name], axes=axes, keepdims=keepdims)
            return [reduce_node]
        else:
            create_tensor([1], name+'_1', kwargs['initializer'])
            create_tensor([0], name+'_0', kwargs['initializer'])
            create_tensor([len(axes)], name+'_axes_dim', kwargs['initializer'])
            nodes = [
                make_node(onnx_op_name, input_nodes, [name+'_reduce'], axes=axes, keepdims=keepdims),
                make_node('Shape', [name+'_reduce'], [name+'_reduce_shape']),
                make_node('Shape', [name+'_reduce_shape'], [name+'_reduce_dim']),
                make_node('Shape', [input_nodes[0]], [name+'_in_shape']),
                make_node('Shape', [name+'_in_shape'], [name+'_in_dim']),
                make_node('Equal', [name+'_axes_dim', name+'_in_dim'], [name+'_equal']),
                make_node('Where', [name+'_equal', name+'_1', name+'_reduce_dim'], [name+'_where0']),
                make_node('Tile', [name+'_0', name+'_where0'], [name+'_tile']),
                make_node('Unsqueeze', [name+'_0', name+'_0'], [name+'_unsqueeze']),
                make_node('Where', [name+'_equal', name+'_1', name+'_0'], [name+'_where1']),
                make_node('ScatterND', [name+'_tile', name+'_unsqueeze', name+'_where1'], [name+'_SND']),
                make_node('Reshape', [name+'_reduce', name+'_SND'], [name]),
            ]
            return nodes
    else:

        if keepdims:
            reduce_node = make_node(onnx_op_name, input_nodes, [name], keepdims=keepdims)
            return [reduce_node]
        else:
            create_tensor([1], name+'_1', kwargs['initializer'])
            nodes = [
                make_node(onnx_op_name, input_nodes, [name+'_norm'], keepdims=keepdims),
                make_node('Reshape', [name+'_norm', name+'_1'], [name])
            ]
            return nodes


@mx_op.register("log_softmax", OPSET_VERSION)
def convert_logsoftmax(node, **kwargs):
    """Map MXNet's log_softmax operator attributes to onnx's LogSoftMax operator
    and return the created node.
    """
    name, input_nodes, attrs = get_inputs(node, kwargs)

    # Converting to int
    axis = int(attrs.get("axis", -1))
    temp = attrs.get('temperature', 'None')
    use_length = attrs.get('use_length', 'False')

    if temp != 'None':
        raise AttributeError('LogSoftMax currently does not support temperature!=None')

    if use_length in ['1', 'True']:
        raise AttributeError('LogSoftMax currently does not support use_length==True')

    node = onnx.helper.make_node(
        'LogSoftmax',
        input_nodes,
        [name],
        axis=axis,
        name=name
    )

    return [node]


@mx_op.register('_split_v2', OPSET_VERSION)
def convert_contrib_split_v2(node, **kwargs):
    """Map MXNet's _split_v2 operator
    """
    from onnx.helper import make_node
    name, input_nodes, attrs = get_inputs(node, kwargs)
    axis = int(attrs.get('axis', 0))
    squeeze_axis = attrs.get('squeeze_axis', 'False')
    sections = int(attrs.get('sections', 0))
    indices = convert_string_to_list(attrs.get('indices', '[]'))
    if sections <= 0 and len(indices) == 0:
        raise NotImplementedError('section or indices must be set')
    if sections > 0:
        output_nodes = [name+str(i) for i in range(sections)]
        if squeeze_axis == 'False':
            nodes = [
                make_node('Split', input_nodes, output_nodes, axis=axis),
            ]
        else:
            output_nodes_ = [name+str(i)+'_' for i in range(sections)]
            create_tensor([axis], name+'_axis', kwargs['initializer'])
            nodes = [
                make_node('Split', input_nodes, output_nodes_, axis=axis),
            ]
            for i in range(sections):
                nodes += [
                    make_node("Squeeze", [output_nodes_[i], name+'_axis'], [output_nodes[i]]),
                ]
    else:
        indices.sort()
        split = []
        for i in range(1, len(indices)):
            if indices[i] >= indices[i-1]:
                split.append(indices[i] - indices[i-1])

        output_nodes = [name+str(i) for i in range(len(split)+1)]
        create_tensor([0], name+'_0', kwargs['initializer'])
        create_tensor([axis], name+'_axis', kwargs['initializer'])
        create_tensor([axis+1], name+'_axis+1', kwargs['initializer'])
        create_tensor(split, name+'_split_', kwargs['initializer'])
        create_tensor([sum(split)], name+'_sum', kwargs['initializer'])
        nodes = [
            make_node('Shape', input_nodes, [name+'_shape']),
            make_node('Slice', [name+'_shape', name+'_axis', name+'_axis+1', name+'_0'], [name+'_dim']),
            make_node('Sub', [name+'_dim', name+'_sum'], [name+'_sub']),
            make_node('Concat', [name+'_split_', name+'_sub'], [name+'_concat'], axis=0),
            make_node('Less', [name+'_concat', name+'_0'], [name+'_less']),
            make_node('Where', [name+'_less', name+'_0', name+'_concat'], [name+'_split']),
            ]
        if squeeze_axis == 'False':
            nodes += [
                make_node('Split', [input_nodes[0], name+'_split'], output_nodes, axis=axis),
            ]
        else:
            output_nodes_ = [name+str(i)+'_' for i in range(len(split)+1)]
            nodes += [
                make_node('Split', [input_nodes[0], name+'_split'], output_nodes_, axis=axis),
            ]
            for i, output_node in enumerate(output_nodes):
                nodes += [
                    make_node("Squeeze", [output_nodes_[i], name+'_axis'], [output_node]),
                ]

    return nodes
