import os
import math
import numpy as np
import six
import megengine._internal as mgb
from enum import Enum

from py_proto import mace_pb2
from transform import base_converter
from transform.base_converter import PoolingType
from transform.base_converter import ActivationType
from transform.base_converter import EltwiseType
from transform.base_converter import FrameworkType
from transform.base_converter import ReduceType
from transform.base_converter import DataFormat
from transform.base_converter import MaceOp
from transform.base_converter import MaceKeyword
from transform.base_converter import ConverterUtil
from transform.base_converter import RoundMode
from utils.util import mace_check

mge_kernel_h_str = "window_h"
mge_kernel_w_str = "window_w"
mge_stride_h_str = "stride_h"
mge_stride_w_str = "stride_w"
mge_pad_h_str = "pad_h"
mge_pad_w_str = "pad_w"
mge_dilate_h_str = "dilate_h"
mge_dilate_w_str = "dilate_w"

MGESupportedOps = [
    "AxisAddRemove",
    "BatchNormForward",
    "Concat",
    "ConvolutionForward",
    "ConvolutionBackwardData",
    "Dimshuffle",
    "Elemwise",
    "GetVarShape",
    "Host2DeviceCopy",
    "Identity",
    "MarkNoBroadcastElemwise",
    "MatrixMul",
    "PoolingForward",
    "Reduce",
    "Reshape",
    "SharedDeviceTensor",
    "Subtensor",
]

MGEOpType = Enum("MGEOpType", [(op, op) for op in MGESupportedOps], type=str)


def get_symvar_value(mge_symvar):
    if mge_symvar.inferred_value is not None:
        val = mge_symvar.inferred_value
    else:
        cg = mge_symvar.owner_graph
        func = cg.compile_outonly(mge_symvar)
        val = func()

    return val


def is_consumer_group_conv(mge_symvar, var2oprs, map_oprs):
    consumer_ids = var2oprs[mge_symvar.id]
    n_consumers = len(consumer_ids)
    for consumer_id in consumer_ids:
        consumer_op = map_oprs[consumer_id[0]]
        if (mgb.cgtools.get_opr_type(consumer_op)
                in ("ConvolutionForward", "ConvolutionBackwardData")
                and consumer_op.params["sparse"] == "GROUP"):
            mace_check(n_consumers == 1,
                       "This tensor should only feed depthwise conv/deconv")
            return True
    return False


class MegengineConverter(base_converter.ConverterInterface):
    """A class for convert megengine dumped model to mace model."""

    compute_format_type = {
        "NCHW": DataFormat.NCHW,
        "NHWC": DataFormat.NHWC,
        "DEFAULT": DataFormat.NCHW,
    }

    reduce_math_type = {
        "SUM": ReduceType.SUM,
        "PROD": ReduceType.PROD,
        "MIN": ReduceType.MIN,
        "MAX": ReduceType.MAX,
    }

    # SQE_DIFF, CLIP, SIGN maybe needed
    eltwise_type = {
        "ADD": EltwiseType.SUM,
        "SUB": EltwiseType.SUB,
        "MUL": EltwiseType.PROD,
        "TRUE_DIV": EltwiseType.DIV,
        "MIN": EltwiseType.MIN,
        "MAX": EltwiseType.MAX,
        "NEGATE": EltwiseType.NEG,
        "ABS": EltwiseType.ABS,
        "POW": EltwiseType.POW,
        "EQ": EltwiseType.EQUAL,
        "FLOOR_DIV": EltwiseType.FLOOR_DIV,
        "EXP": EltwiseType.POW,
    }

    activation_type = {
        "RELU": ActivationType.RELU,
        "TANH": ActivationType.TANH,
        "SIGMOID": ActivationType.SIGMOID,
    }

    def __init__(self, option, src_model_file):
        self._op_converters = {
            MGEOpType.AxisAddRemove.name: self.convert_axisaddrm,
            MGEOpType.BatchNormForward.name: self.convert_batchnorm,
            MGEOpType.Concat.name: self.convert_concat,
            MGEOpType.ConvolutionForward.name: self.convert_conv2d,
            MGEOpType.ConvolutionBackwardData.name: self.convert_deconv2d,
            MGEOpType.Dimshuffle.name: self.convert_dimshuffle,
            MGEOpType.Elemwise.name: self.convert_elemwise,
            MGEOpType.GetVarShape.name: self.convert_shape,
            MGEOpType.Host2DeviceCopy.name: self.convert_nop,
            MGEOpType.Identity.name: self.convert_identity,
            MGEOpType.MarkNoBroadcastElemwise.name: self.convert_identity,
            MGEOpType.MatrixMul.name: self.convert_matmul,
            MGEOpType.PoolingForward.name: self.convert_pooling,
            MGEOpType.Reduce.name: self.convert_reduce,
            MGEOpType.Reshape.name: self.convert_reshape,
            MGEOpType.SharedDeviceTensor.name: self.convert_nop,
            MGEOpType.Subtensor.name: self.convert_subtensor,
        }
        self._option = option
        self._mace_net_def = mace_pb2.NetDef()
        ConverterUtil.set_filter_format(self._mace_net_def, DataFormat.OIHW)
        ConverterUtil.add_data_format_arg(self._mace_net_def, DataFormat.NCHW)

        cg, _, outputs = mgb.load_comp_graph_from_file(src_model_file)
        map_oprs, _, var2oprs, *_ = mgb.cgtools.graph_traversal(outputs)
        # prune second input of reshape
        # because it introduces several ops, may increase the overhead
        operators = mgb.cgtools.get_oprs_seq(outputs, prune_reshape=True)

        self._mge_cg = cg
        self._mge_operators = operators
        self._mge_map_oprs = map_oprs
        self._mge_var2oprs = var2oprs

        self._skip_tensors = set()
        self._bn_statistis_tensors = {}

    def run(self):
        self.convert_ops()

        self.replace_input_output_tensor_name()
        return self._mace_net_def

    # only change the input/output tensor name for whole model
    def replace_input_output_tensor_name(self):
        for op in self._mace_net_def.op:
            for i in six.moves.range(len(op.input)):
                if "," in op.input[i]:
                    op_name = op.input[i]
                    op_name = op_name.replace(",", "#")
                    if (op_name in self._option.input_nodes or
                            op_name in self._option.output_nodes):
                        op.input[i] = op_name
            for i in six.moves.range(len(op.output)):
                if "," in op.output[i]:
                    op_name = op.output[i]
                    op_name = op_name.replace(",", "#")
                    if op_name in self._option.output_nodes:
                        op.output[i] = op_name

    # this method will be called by convert_conv2d/deconv2d and convert_pooling
    @staticmethod
    def add_stride_pad_kernel_arg(params, op_def):
        stride = [params[mge_stride_h_str], params[mge_stride_w_str]]
        pad = [params[mge_pad_h_str] * 2, params[mge_pad_w_str] * 2]

        strides_arg = op_def.arg.add()
        strides_arg.name = MaceKeyword.mace_strides_str
        strides_arg.ints.extend(stride)
        padding_arg = op_def.arg.add()
        padding_arg.name = MaceKeyword.mace_padding_values_str
        padding_arg.ints.extend(pad)

        if op_def.type == MaceOp.Pooling.name:
            kernel = [params[mge_kernel_h_str], params[mge_kernel_w_str]]
            kernels_arg = op_def.arg.add()
            kernels_arg.name = MaceKeyword.mace_kernel_str
            kernels_arg.ints.extend(kernel)
        if op_def.type in (MaceOp.Conv2D.name, MaceOp.DepthwiseConv2d.name,
                           MaceOp.Deconv2D.name,
                           MaceOp.DepthwiseDeconv2d.name):
            dilation = [params[mge_dilate_h_str], params[mge_dilate_w_str]]
            dilation_arg = op_def.arg.add()
            dilation_arg.name = MaceKeyword.mace_dilations_str
            dilation_arg.ints.extend(dilation)

    def convert_ops(self):
        for mge_op in self._mge_operators:
            opr_type = mgb.cgtools.get_opr_type(mge_op)

            # some reshape operators provide data for batchnorm
            if opr_type == "Reshape":
                output = mge_op.outputs[0]
                next_ops = self._mge_var2oprs[output.id]
                if len(next_ops) == 1:
                    (next_op_id, _) = next_ops[0]
                    next_op = self._mge_map_oprs[next_op_id]

                    if mgb.cgtools.get_opr_type(next_op) == "BatchNormForward":
                        self._skip_tensors.update(
                            [inp.name for inp in mge_op.inputs])
                        # using output name to address input symbol var
                        self._bn_statistis_tensors[mge_op.outputs[0].name] = \
                            mge_op.inputs[0]
                        # skip this reshape op
                        continue

            self._op_converters[opr_type](mge_op)

        self.convert_tensors()

    def add_tensor(self, name, shape, data_type, value):
        tensor = self._mace_net_def.tensors.add()
        tensor.name = name
        tensor.dims.extend(list(shape))
        tensor.data_type = data_type
        if data_type == mace_pb2.DT_INT32:
            tensor.int32_data.extend(value)
        else:
            tensor.float_data.extend(value)

    # convert all pre-calculated and constant tensors
    def convert_tensors(self):
        for mge_op in self._mge_operators:
            type_opr = mgb.cgtools.get_opr_type(mge_op)

            # all tensors generated by SharedDeviceTensor op
            if type_opr == "SharedDeviceTensor":
                output = mge_op.outputs[0]
                if output.name not in self._skip_tensors:
                    nshape = output.imm_shape
                    # tensor used for depthwise conv/deconv should be reshaped
                    for_group_conv = is_consumer_group_conv(
                        output, self._mge_var2oprs, self._mge_map_oprs
                    )
                    if for_group_conv:
                        nshape = (
                            1,
                            output.imm_shape[0],
                            output.imm_shape[3],
                            output.imm_shape[4],
                        )

                    self.add_tensor(
                        output.name,
                        nshape,
                        mace_pb2.DT_FLOAT,
                        get_symvar_value(output).flatten())
            else:
                # handle all constant values
                for const_tensor in mge_op.inputs:
                    if (const_tensor.inferred_value is not None
                            and const_tensor.name not in self._skip_tensors):
                        self.add_tensor(
                            const_tensor.name,
                            const_tensor.imm_shape,
                            mace_pb2.DT_INT32,
                            const_tensor.inferred_value.flatten())

    def convert_nop(self, mge_op):
        pass

    def convert_general_op(self, mge_op):
        op = self._mace_net_def.op.add()
        op.name = mge_op.name
        op.type = mgb.cgtools.get_opr_type(mge_op)
        op.input.extend([mge_input.name for mge_input in mge_op.inputs])
        op.output.extend([mge_output.name for mge_output in mge_op.outputs])
        for mge_output in mge_op.outputs:
            output_shape = op.output_shape.add()
            output_shape.dims.extend(mge_output.imm_shape)

        data_type_arg = op.arg.add()
        data_type_arg.name = "T"
        data_type_arg.i = self._option.data_type

        framework_type_arg = op.arg.add()
        framework_type_arg.name = MaceKeyword.mace_framework_type_str
        framework_type_arg.i = FrameworkType.MEGENGINE.value

        # check compute format of megengine
        compute_format = DataFormat.NCHW
        try:
            if "format" in mge_op.params.keys():
                compute_format = self.compute_format_type[
                    mge_op.params["format"]
                ]
        except AttributeError:
            compute_format = DataFormat.NCHW
        ConverterUtil.add_data_format_arg(op, compute_format)

        return op

    def convert_identity(self, mge_op):
        op = self.convert_general_op(mge_op)
        op.type = MaceOp.Identity.name

    def convert_conv2d(self, mge_op):
        op = self.convert_general_op(mge_op)

        if mge_op.params["sparse"] == "GROUP":
            # weight shape in group conv2d:
            # (groups, out_channel//groups, in_channels//groups, *kernel_size)
            groups_divisible = mge_op.inputs[1].imm_shape[2]
            mace_check(
                groups_divisible == 1,
                "Mace does not support group convolution yet",
            )
            op.type = MaceOp.DepthwiseConv2d.name
        elif mge_op.params["sparse"] == "DENSE":
            op.type = MaceOp.Conv2D.name
        else:
            raise Exception("Unknown sparse mode")

        mace_check(
            mge_op.params["mode"] != "CONVOLUTION",
            "Mace does not support CONVOLUTION computation mode yet",
        )

        self.add_stride_pad_kernel_arg(mge_op.params, op)

        del op.output[1:]
        del op.output_shape[1:]

    def convert_deconv2d(self, mge_op):
        op = self.convert_general_op(mge_op)

        if mge_op.params["sparse"] == "GROUP":
            # weight shape in group conv2d:
            # (groups, out_channel//groups, in_channels//groups, *kernel_size)
            groups_divisible = mge_op.inputs[0].imm_shape[2]
            mace_check(
                groups_divisible == 1,
                "Mace does not support group deconvolution yet",
            )
            op.type = MaceOp.DepthwiseConv2d.name
        elif mge_op.params["sparse"] == "DENSE":
            op.type = MaceOp.Deconv2D.name
        else:
            mace_check(False, "Unknown sparse mode")

        mace_check(
            mge_op.params["mode"] != "CONVOLUTION",
            "Mace does not support CONVOLUTION computation mode yet",
        )

        self.add_stride_pad_kernel_arg(mge_op.params, op)

        # inputs order is strange in megengine, fix it
        swaped_list = [op.input[1], op.input[0]]
        del op.input[:]
        op.input.extend(swaped_list)

        del op.output[1:]
        del op.output_shape[1:]

    def convert_dimshuffle(self, mge_op):
        op = self.convert_general_op(mge_op)
        op.type = MaceOp.Transpose.name

        dims_arg = op.arg.add()
        dims_arg.name = MaceKeyword.mace_dims_str
        dims_arg.ints.extend(mge_op.params["pattern"])

    def convert_math_elemwise(self, mge_op):
        op = self.convert_general_op(mge_op)
        op.type = MaceOp.Eltwise.name

        type_arg = op.arg.add()
        type_arg.name = MaceKeyword.mace_element_type_str
        type_arg.i = self.eltwise_type[mge_op.params["mode"]].value
        # EXP in megengine always use the np.e as base
        if mge_op.params["mode"] == "EXP":
            exp_tensor_name = mge_op.name + "_exp_base"
            exp_shape = mge_op.outputs[0].imm_shape
            exp_value = (np.e * np.ones(exp_shape)).flatten()
            self.add_tensor(
                exp_tensor_name, exp_shape, mace_pb2.DT_FLOAT, exp_value
            )
            del op.input[0]
            op.input.extend([exp_tensor_name, mge_op.inputs[0].name])

    def convert_activation(self, mge_op):
        op = self.convert_general_op(mge_op)
        op.type = MaceOp.Activation.name

        type_arg = op.arg.add()
        type_arg.name = MaceKeyword.mace_activation_type_str
        type_arg.s = six.b(self.activation_type[mge_op.params["mode"]].name)

    def convert_elemwise(self, mge_op):
        mode = mge_op.params["mode"]
        if mode in self.eltwise_type:
            self.convert_math_elemwise(mge_op)
        else:
            self.convert_activation(mge_op)

    def convert_pooling(self, mge_op):
        op = self.convert_general_op(mge_op)
        op.type = MaceOp.Pooling.name

        pool_type_arg = op.arg.add()
        pool_type_arg.name = MaceKeyword.mace_pooling_type_str

        round_mode_arg = op.arg.add()
        round_mode_arg.name = MaceKeyword.mace_round_mode_str
        round_mode_arg.i = RoundMode.FLOOR.value

        # check the case of counting include padding
        mode = mge_op.params["mode"]
        if mode == "AVERAGE_COUNT_EXCLUDE_PADDING" or \
                (mode == "AVERAGE" and mge_op.params["pad_w"] == 0 and
                 mge_op.params["pad_h"] == 0):
            pool_type_arg.i = PoolingType.AVG.value
        elif mode == "MAX":
            pool_type_arg.i = PoolingType.MAX.value
        else:
            mace_check(False,
                       "AVERAGE pooling should not count padding values")

        self.add_stride_pad_kernel_arg(mge_op.params, op)

        # delete workspace output, it's useless
        del op.output[1:]
        del op.output_shape[1:]

    def convert_matmul(self, mge_op):
        op = self.convert_general_op(mge_op)
        op.type = MaceOp.MatMul.name

        transpose_a = mge_op.params["transposeA"]
        transpose_a_arg = op.arg.add()
        transpose_a_arg.name = MaceKeyword.mace_transpose_a_str
        transpose_a_arg.i = int(transpose_a)

        transpose_b = mge_op.params["transposeB"]
        transpose_b_arg = op.arg.add()
        transpose_b_arg.name = MaceKeyword.mace_transpose_b_str
        transpose_b_arg.i = int(transpose_b)

        del op.output[1:]
        del op.output_shape[1:]

    def convert_reshape(self, mge_op):
        op = self.convert_general_op(mge_op)
        op.type = MaceOp.Reshape.name

        # just use the output shape
        del op.input[1]
        t_shape = list(mge_op.outputs[0].imm_shape)
        shape_tensor_name = mge_op.name + "_dest_shape"
        self.add_tensor(
            shape_tensor_name, [len(t_shape)], mace_pb2.DT_INT32, t_shape
        )
        op.input.extend([shape_tensor_name])

    # usually after reduce operator, remove dimension with value 1
    # it's hard to just follow this operator
    # sometimes axis-add and axis-remove may exist at the same time
    # for complicated use-case, using reshape is easier
    def convert_axisaddrm(self, mge_op):
        op = self.convert_general_op(mge_op)
        if mge_op.params["nr_desc"] == 1:
            if mge_op.params["desc"][0]["method"] == 0:
                op.type = MaceOp.ExpandDims.name
            else:
                op.type = MaceOp.Squeeze.name

            axis_arg = op.arg.add()
            axis_arg.name = MaceKeyword.mace_axis_str
            axis_arg.i = mge_op.params["desc"][0]["axisnum"]
        else:
            op.type = MaceOp.Reshape.name

            dest_shape_tensor_name = op.name + "_dest_shape"
            dest_shape = mge_op.outputs[0].imm_shape
            self.add_tensor(
                dest_shape_tensor_name,
                (len(dest_shape),),
                mace_pb2.DT_INT32,
                dest_shape,
            )
            op.input.extend([dest_shape_tensor_name])

    def convert_reduce(self, mge_op):
        op = self.convert_general_op(mge_op)
        op.type = MaceOp.Reduce.name

        reduce_type_arg = op.arg.add()
        reduce_type_arg.name = MaceKeyword.mace_reduce_type_str
        reduce_type_arg.i = self.reduce_math_type[mge_op.params["mode"]].value

        # in megengine axis won't be list, just int
        axis_arg = op.arg.add()
        axis_arg.name = MaceKeyword.mace_axis_str
        axis_arg.ints.append(mge_op.params["axis"])

        # megengine will always keep dims in Reduce operator
        # dim removal will be done by operator AxisAddRemove
        keep_dims_arg = op.arg.add()
        keep_dims_arg.name = MaceKeyword.mace_keepdims_str
        keep_dims_arg.i = 1

        del op.output[1:]
        del op.output_shape[1:]

    def convert_concat(self, mge_op):
        op = self.convert_general_op(mge_op)
        op.type = MaceOp.Concat.name

        axis_arg = op.arg.add()
        axis_arg.name = MaceKeyword.mace_axis_str
        axis_arg.i = mge_op.params["axis"]

    def convert_batchnorm(self, mge_op):
        op = self.convert_general_op(mge_op)
        op.type = MaceOp.BatchNorm.name

        gamma_value = get_symvar_value(
            self._bn_statistis_tensors[mge_op.inputs[1].name]
        ).flatten()
        beta_value = get_symvar_value(
            self._bn_statistis_tensors[mge_op.inputs[2].name]
        ).flatten()
        mean_value = get_symvar_value(mge_op.inputs[3]).flatten()
        var_value = get_symvar_value(mge_op.inputs[4]).flatten()
        epsilon_value = 1e-5

        scale_name = mge_op.name + "_scale"
        offset_name = mge_op.name + "_offset"
        scale_value = (1.0 / np.vectorize(math.sqrt)(
            var_value + epsilon_value)) * gamma_value
        offset_value = (-mean_value * scale_value) + beta_value
        self.add_tensor(
            scale_name, scale_value.shape, mace_pb2.DT_FLOAT, scale_value
        )
        self.add_tensor(
            offset_name, offset_value.shape, mace_pb2.DT_FLOAT, offset_value
        )
        self._skip_tensors.update([inp.name for inp in mge_op.inputs][1:])

        del op.input[1:]
        op.input.extend([scale_name, offset_name])
        # outputs[4] is the correct output
        del op.output[-1:]
        del op.output_shape[-1:]
        del op.output[:4]
        del op.output_shape[:4]

    def convert_shape(self, mge_op):
        op = self.convert_general_op(mge_op)
        op.type = MaceOp.Shape.name
        op.output_type.extend([mace_pb2.DT_INT32])

    # axis of subtensor should be constant
    # subtensor in megengine: numpy-like indexing
    def convert_subtensor(self, mge_op):
        op1 = self.convert_general_op(mge_op)
        op1.type = MaceOp.StridedSlice.name

        axis = mge_op.inputs[1].inferred_value
        t_shape = list(mge_op.inputs[0].imm_shape)

        begin_tensor_name = mge_op.name + "_begin"
        end_tensor_name = mge_op.name + "_end"
        stride_tensor_name = mge_op.name + "_stride"
        begin_tensor_shape = (len(t_shape),)
        end_tensor_shape = (len(t_shape),)
        stride_tensor_shape = (len(t_shape),)

        begin_vals = [0] * len(t_shape)
        end_vals = [shapei for shapei in t_shape]
        stride_vals = [1] * len(t_shape)

        def check_val(sym_var):
            try:
                val = sym_var.inferred_value[0]
            except TypeError:
                mace_check(
                    False, "you should feed const values for subtensor axis"
                )
            return val

        squeeze_dims = []
        idx = len(mge_op.inputs) - 1
        while idx:
            val = check_val(mge_op.inputs[idx])
            for ai in mge_op.params[::-1]:
                ai_idx = ai["axis"]
                if ai["step"] > 0:
                    stride_vals[ai_idx] = val
                    idx -= 1
                    if idx == 0:
                        break
                    val = check_val(mge_op.inputs[idx])
                if ai["end"] > 0:
                    if val < 0:
                        val = t_shape[ai_idx] + val
                    end_vals[ai_idx] = val
                    idx -= 1
                    if idx == 0:
                        break
                    val = check_val(mge_op.inputs[idx])
                if ai["begin"] > 0:
                    if val < 0:
                        val = t_shape[ai_idx] + val
                    begin_vals[ai_idx] = val
                    idx -= 1
                    if idx == 0:
                        break
                    val = check_val(mge_op.inputs[idx])
                if ai["idx"] > 0:
                    if val < 0:
                        val = t_shape[ai_idx] + val
                    squeeze_dims.append(ai_idx)
                    begin_vals[ai_idx] = val
                    end_vals[ai_idx] = val + 1
                    idx -= 1
                    if idx == 0:
                        break
                    val = check_val(mge_op.inputs[idx])

        for ai_idx in range(len(t_shape)):
            t_shape[ai_idx] = math.ceil(
                (end_vals[ai_idx] - begin_vals[ai_idx]) / stride_vals[ai_idx]
            )

        self.add_tensor(
            begin_tensor_name,
            begin_tensor_shape,
            mace_pb2.DT_INT32,
            begin_vals,
        )
        self.add_tensor(
            end_tensor_name, end_tensor_shape, mace_pb2.DT_INT32, end_vals
        )
        self.add_tensor(
            stride_tensor_name,
            stride_tensor_shape,
            mace_pb2.DT_INT32,
            stride_vals,
        )

        del op1.input[1:]
        op1.input.extend(
            [begin_tensor_name, end_tensor_name, stride_tensor_name]
        )

        if len(squeeze_dims) > 0:
            # create squeeze op to remove shape=1 dims
            mid_output_name = mge_op.name + "_mid_reshape"

            del op1.output[0]
            op1.output.extend([mid_output_name])
            output_shape = op1.output_shape[0]
            del output_shape.dims[:]
            output_shape.dims.extend(t_shape)

            op2 = self._mace_net_def.op.add()
            op2.type = MaceOp.Squeeze.name
            op2.name = mge_op.name + "_squeeze"

            data_type_arg = op2.arg.add()
            data_type_arg.name = "T"
            data_type_arg.i = self._option.data_type

            framework_type_arg = op2.arg.add()
            framework_type_arg.name = MaceKeyword.mace_framework_type_str
            framework_type_arg.i = FrameworkType.MEGENGINE.value

            ConverterUtil.add_data_format_arg(op2, DataFormat.NCHW)

            op2.input.extend([mid_output_name])
            op2.output.extend([mge_op.outputs[0].name])
            output_shape = op2.output_shape.add()
            output_shape.dims.extend(mge_op.outputs[0].imm_shape)

            axis_arg = op2.arg.add()
            axis_arg.name = MaceKeyword.mace_axis_str
            axis_arg.ints.extend(squeeze_dims)
