base.py

from abc import ABC, abstractmethod
from distutils.version import LooseVersion

import inspect
import math
import warnings
import torch
import typing

import numpy as np

from .. import tflite as tfl

from tinynn.util.util import get_logger

log = get_logger(__name__, 'INFO')


class OperatorConverter(ABC):
    def __init__(
        self,
        node,
        tensor_map,
        scope_name,
        asymmetric=True,
        q_type=np.uint8,
        hybrid_q_type=np.int8,
        map_bilstm_to_lstm=False,
        enable_mtk_ops=False,
        hybrid_asymmetric_inputs=False,
        unroll_rnn=False,
        separated_rnn_gate_calc=False,
        conv_transpose_with_bias=True,
        legacy_gelu=False,
    ) -> None:
        self.scope_name = scope_name
        self.input_names = self.get_input_names(node)
        self.output_names = self.get_output_names(node)
        self.input_tensors = self.get_input_tensors(tensor_map)
        self.output_tensors = []
        self.output_nodes = []
        self.ops = []
        self.attr_count = 0
        self.transform_count = 0
        self.asymmetric = asymmetric
        self.q_type = q_type
        self.hybrid_q_type = hybrid_q_type
        self.map_bilstm_to_lstm = map_bilstm_to_lstm
        self.enable_mtk_ops = enable_mtk_ops
        self.hybrid_asymmetric_inputs = hybrid_asymmetric_inputs
        self.unroll_rnn = unroll_rnn
        self.separated_rnn_gate_calc = separated_rnn_gate_calc
        self.conv_transpose_with_bias = conv_transpose_with_bias
        self.legacy_gelu = legacy_gelu

    @abstractmethod
    def parse(self, node, attrs, args, graph_converter):
        pass

    def get_tensor_name(self, tensor_name, scope_name=None):
        if scope_name is None:
            scope_name = self.scope_name

        if scope_name:
            return f'{scope_name}_{tensor_name}'
        else:
            return tensor_name

    def get_input_names(self, node):
        return [self.get_tensor_name(x.debugName()) for x in list(node.inputs())]

    def get_output_names(self, node):
        return [self.get_tensor_name(x.debugName()) for x in list(node.outputs())]

    def get_input_tensors(self, tensor_map):
        input_tensors = []
        for n in self.input_names:
            if n in tensor_map:
                input_tensors.append(tensor_map[n])
            else:
                raise Exception(f'{n} is not found in the tensor map')
        return input_tensors

    def get_output_tensors(self):
        return self.output_tensors

    def get_ops(self):
        return self.ops

    @staticmethod
    def fetch_all_attrs(node):
        attrs = {}
        for name in node.attributeNames():
            attrs[name] = get_prop_from_node(node, name, return_type=True)
        return attrs

    def fetch_annotated_args(self, node):
        if len(self.input_tensors) == 0:
            return dict()

        k = node.kind()
        if k.startswith('prim::'):
            return dict()

        schemas = torch._C._jit_get_schemas_for_operator(k)
        candidates = []
        for schema in schemas:
            if 'name' in schema.overload_name:
                continue
            if len(schema.arguments) == len(self.input_tensors):
                candidates.append(schema)

        assert len(candidates) > 0, f"Cannot find the schema for {k}({self.output_names[0]})"
        names = (x.name for x in candidates[0].arguments)

        # TODO: Better selection for multiple schemas
        return dict(zip(names, range(len(self.input_tensors))))

    def unimplemented(self, node, attrs, args):
        log.debug(f'node: {node}')
        log.debug('inputs:')
        for name, tensors in zip(self.input_names, self.input_tensors):
            if type(tensors) not in (list, tuple):
                tensors = [tensors]
            for tensor in tensors:
                log.debug(f'name: {name}')
                log.debug(f'tensor: {tensor}')
                if hasattr(tensor, 'shape'):
                    log.debug(f'shape: {tensor.shape}')
                if hasattr(tensor, 'dtype'):
                    log.debug(f'dtype: {tensor.dtype}')
                log.debug('-' * 60)
        log.debug('outputs:')
        for name, tensors in zip(self.output_names, self.output_tensors):
            if type(tensors) not in (list, tuple):
                tensors = [tensors]
            for tensor in tensors:
                log.debug(f'name: {name}')
                log.debug(f'tensor: {tensor}')
                if hasattr(tensor, 'shape'):
                    log.debug(f'shape: {tensor.shape}')
                if hasattr(tensor, 'dtype'):
                    log.debug(f'dtype: {tensor.dtype}')
                log.debug('-' * 60)
        log.debug(f'attrs: {attrs}')
        log.debug(f'args: {args}')
        raise NotImplementedError

    def run(self, node):
        kind = node.kind()
        inplace = kind.endswith('_')
        func = torch._C._jit_get_operation(kind)
        if inplace:
            tmp_inputs = [x.detach().clone() if isinstance(x, torch.Tensor) else x for x in self.input_tensors]
        if isinstance(func, tuple):
            func = func[0]
        with torch.no_grad():
            legacy = True
            if LooseVersion(torch.__version__) >= LooseVersion('1.8.0'):
                try:
                    o = func(*self.input_tensors)
                    legacy = False
                except (TypeError, RuntimeError):
                    pass

            if legacy:
                try:
                    args = self.fetch_annotated_args(node)
                    kwargs = dict(zip(args.keys(), self.input_tensors))
                    o = func(**kwargs)
                except RuntimeError as e:
                    if 'device' in kwargs:
                        kwargs['device'] = 0
                        o = func(**kwargs)
                    else:
                        raise e

        if inplace:
            self.input_tensors.clear()
            self.input_tensors.extend(tmp_inputs)

        if len(self.output_names) == 1:
            self.output_tensors.append(o)
        else:
            self.output_tensors.extend(o)

    def to_tfl_tensors(
        self, names, tensors, has_buffers=None, graph_converter=None, non_existent_as_buffer=False
    ) -> typing.List[tfl.Tensor]:
        tfl_tensors = []
        if has_buffers is None:
            has_buffers = [None] * len(tensors)
        elif type(has_buffers) is bool:
            has_buffers = [has_buffers] * len(tensors)
        assert len(names) == len(tensors) == len(has_buffers)
        for n, t, b in zip(names, tensors, has_buffers):
            if b is None:
                if graph_converter is not None and n in graph_converter.tensor_map:
                    t = graph_converter.tensor_map[n]
                else:
                    t = tfl.Tensor(
                        t, n, has_buffer=non_existent_as_buffer, asymmetric=self.asymmetric, q_type=self.q_type
                    )
            else:
                t = tfl.Tensor(t, n, has_buffer=b, asymmetric=self.asymmetric, q_type=self.q_type)
            tfl_tensors.append(t)
        return tfl_tensors

    def find_or_create_input(self, idx, graph_converter):
        name = self.input_names[idx]
        if name in graph_converter.tensor_map:
            return graph_converter.tensor_map[name]

        # assert has_buffer, 'only tensors with has_buffer=True can be created at this time,' + \
        #     ' when you encounter this message, it means some ops in the computation graph is not supported'''

        tensor = self.input_tensors[idx]
        return tfl.Tensor(tensor, name, has_buffer=True, asymmetric=self.asymmetric, q_type=self.q_type)

    def get_unique_attr_name(self):
        if self.attr_count == 0:
            name = self.output_names[0] + '_attr'
        else:
            name = self.output_names[0] + f'_attr_{self.attr_count}'
        self.attr_count += 1
        return name

    def get_unique_transform_name(self):
        if self.transform_count == 0:
            name = self.output_names[0] + '_transform'
        else:
            name = self.output_names[0] + f'_transform_{self.transform_count}'
        self.transform_count += 1
        return name

    def create_transform_tensor(self, tensor, name=None, quantization=None):
        if name is None:
            name = self.get_unique_transform_name()
        return tfl.Tensor(
            tensor, name, has_buffer=False, quantization=quantization, asymmetric=self.asymmetric, q_type=self.q_type
        )

    def create_attr_tensor(self, tensor, name=None, hybrid=False, quantization=None):
        if name is None:
            name = self.get_unique_attr_name()

        if hybrid:
            q_type = np.int8
        else:
            q_type = self.q_type

        tensor = tfl.Tensor(
            tensor, name, has_buffer=True, quantization=quantization, asymmetric=self.asymmetric, q_type=q_type
        )

        if hybrid and self.hybrid_q_type == np.uint8:
            tensor.reinterpret_as(self.hybrid_q_type)

        return tensor

    def unpack_params(self, params):
        result = {}
        for method in params._method_names():
            if not (method.startswith('__') and method.endswith('__')):
                result[method] = getattr(params, method)()
        state = params.__getstate__()
        return result, state

    def rescale_weight_scale_for_qnnpack(
        self, input_tensor: tfl.Tensor, weight_tensor: tfl.Tensor, output_tensor: tfl.Tensor
    ):
        updated = False
        orig_scale = weight_tensor.quantization.scale
        while True:
            input_product_scale = input_tensor.quantization.scale * weight_tensor.quantization.scale
            scale = input_product_scale / output_tensor.quantization.scale
            shift = 127 + 31 - 32 - (fp32_to_bits(scale) >> 23)
            if shift >= 32:
                updated = True
                weight_tensor.quantization.scale *= 10
            else:
                break
        if updated:
            cur_scale = weight_tensor.quantization.scale
            log.info(f'rescale quantized weight of {weight_tensor.name}: {orig_scale:.8f}->{cur_scale:.8f}')

    def quantize_numpy(self, tensor, scale, zero_point, dtype=np.uint8):
        q_tensor = np.rint(tensor / scale + zero_point)
        type_info = np.iinfo(dtype)
        if np.any(q_tensor > type_info.max):
            warnings.warn('Overflow while quantizing the tensor')
            q_tensor = np.minimum(q_tensor, type_info.max)
        if np.any(q_tensor < type_info.min):
            warnings.warn('Underflow while quantizing the tensor')
            q_tensor = np.maximum(q_tensor, type_info.min)
        q_tensor = q_tensor.astype(dtype)
        return q_tensor

    def quantize(self, tensor, scale, zero_point, dtype=torch.uint8, dim=None):
        if isinstance(scale, list):
            scale = torch.tensor(scale)
        if isinstance(zero_point, list):
            zero_point = torch.tensor(zero_point)
        q_tensor = torch.round(tensor.detach() / scale + zero_point)
        type_info = torch.iinfo(dtype)
        if (q_tensor > type_info.max).any():
            warnings.warn('Overflow while quantizing the tensor')
            q_tensor[q_tensor > type_info.max] = type_info.max
        if (q_tensor < type_info.min).any():
            warnings.warn('Underflow while quantizing the tensor')
            q_tensor[q_tensor < type_info.min] = type_info.min
        q_tensor = q_tensor.to(dtype=dtype)
        if isinstance(scale, torch.Tensor):
            scale = scale.tolist()
        if isinstance(zero_point, torch.Tensor):
            zero_point = zero_point.tolist()
        return tfl.FakeQuantTensor(q_tensor, scale, zero_point, dim)

    def passthrough(self, graph_converter):
        assert len(self.input_tensors) >= len(self.output_tensors)

        for i in range(len(self.output_tensors)):
            input_tensor = self.input_tensors[i]
            inputs = [self.find_or_create_input(i, graph_converter), self.create_attr_tensor(input_tensor.shape)]
            outputs = self.to_tfl_tensors(self.output_names, self.output_tensors)

            graph_converter.add_operator(tfl.ReshapeOperator(inputs, outputs, input_tensor.shape))

    def elementwise_unary(self, converter_class, graph_converter, *args, **kwargs):
        inputs = [self.find_or_create_input(0, graph_converter)]
        outputs = self.to_tfl_tensors(self.output_names, self.output_tensors)

        if inputs[0].buffer is None:
            graph_converter.add_operator(converter_class(inputs, outputs, *args, **kwargs))

    def elementwise_binary(self, converter_class, graph_converter, autocast, *args, **kwargs):
        if autocast:
            result_dtype = torch.promote_types(self.input_tensors[0].dtype, self.input_tensors[1].dtype)
            for i in range(2):
                t = self.input_tensors[i]
                if result_dtype != t.dtype:
                    casted = t.clone().to(dtype=result_dtype)
                    inp_t = self.find_or_create_input(i, graph_converter)
                    if inp_t.buffer is None:
                        new_inp = self.create_transform_tensor(casted)
                        graph_converter.add_operator(
                            tfl.CastOperator(
                                [inp_t],
                                [new_inp],
                                tfl.torch_tflite_dtype_mappings[t.dtype],
                                tfl.torch_tflite_dtype_mappings[result_dtype],
                            )
                        )
                        self.input_names[i] = new_inp.name
                    self.input_tensors[i] = casted

        inputs = [self.find_or_create_input(i, graph_converter) for i in range(2)]
        if not all((t.buffer is not None for t in inputs)):
            outputs = self.to_tfl_tensors(self.output_names, self.output_tensors)

            graph_converter.add_operator(converter_class(inputs, outputs, *args, **kwargs))

    def reshape(self, graph_converter):
        new_shape = np.array(self.output_tensors[0].shape, dtype='int32')
        inputs = [self.find_or_create_input(0, graph_converter), self.create_attr_tensor(new_shape)]
        outputs = self.to_tfl_tensors(self.output_names, self.output_tensors)

        graph_converter.add_operator(tfl.ReshapeOperator(inputs, outputs, new_shape))

    def wrap_ops_with_dequant_quants(
        self, ops: typing.List[tfl.BaseOperator], input_idx: int = 0, output_idx: int = 0
    ) -> typing.List[tfl.BaseOperator]:
        orig_input = ops[0].inputs[input_idx]
        orig_output = ops[-1].outputs[output_idx]

        new_input = self.create_transform_tensor(orig_input.tensor.astype('float32'))
        new_output = self.create_transform_tensor(orig_output.tensor.astype('float32'))

        dequant_op = tfl.DequantizeOperator([orig_input], [new_input])
        quant_op = tfl.QuantizeOperator([new_output], [orig_output])

        ops[0].inputs[input_idx] = new_input
        ops[-1].outputs[output_idx] = new_output

        return [dequant_op] + ops + [quant_op]

    def wrap_ops_with_2d_3d_reshapes(
        self, ops: typing.List[tfl.BaseOperator], input_idx: int = 0, output_idx: int = 0
    ) -> typing.List[tfl.BaseOperator]:
        orig_input = ops[0].inputs[input_idx]
        orig_output = ops[-1].outputs[output_idx]

        input_shape = np.array(orig_input.tensor.shape[1:], dtype='int32')
        output_shape = np.array(orig_output.tensor.shape, dtype='int32')

        input_shape_tensor = self.create_attr_tensor(input_shape)
        output_shape_tensor = self.create_attr_tensor(output_shape)

        new_input = self.create_transform_tensor(
            orig_input.tensor.reshape(input_shape), quantization=orig_input.quantization
        )
        new_output = self.create_transform_tensor(
            orig_output.tensor.reshape(output_shape[1:]), quantization=orig_output.quantization
        )

        input_reshape_op = tfl.ReshapeOperator([orig_input, input_shape_tensor], [new_input], input_shape)
        output_reshape_op = tfl.ReshapeOperator([new_output, output_shape_tensor], [orig_output], output_shape)

        input_reshape_op.extra_hints['direction'] = 'up'
        output_reshape_op.extra_hints['direction'] = 'down'

        ops[0].inputs[input_idx] = new_input
        ops[-1].outputs[output_idx] = new_output

        return [input_reshape_op] + ops + [output_reshape_op]

    def wrap_ops_with_nhwc_nchw_transposes(
        self, ops: typing.List[tfl.BaseOperator], input_idx: int = 0, output_idx: int = 0
    ) -> typing.List[tfl.BaseOperator]:
        orig_input = ops[0].inputs[input_idx]
        orig_output = ops[-1].outputs[output_idx]

        nhwc2nchw_perm = np.array([0, 3, 1, 2], dtype='int32')
        nchw2nhwc_perm = np.array([0, 2, 3, 1], dtype='int32')

        nhwc2nchw_perm_tensor = self.create_attr_tensor(nhwc2nchw_perm)
        nchw2nhwc_perm_tensor = self.create_attr_tensor(nchw2nhwc_perm)

        new_input = self.create_transform_tensor(
            np.transpose(orig_input.tensor, nchw2nhwc_perm), quantization=orig_input.quantization
        )
        new_output = self.create_transform_tensor(
            np.transpose(orig_output.tensor, nchw2nhwc_perm), quantization=orig_output.quantization
        )

        nchw2nhwc_transpose = tfl.TransposeOperator([orig_input, nchw2nhwc_perm_tensor], [new_input])
        nhwc2nchw_transpose = tfl.TransposeOperator([new_output, nhwc2nchw_perm_tensor], [orig_output])

        nchw2nhwc_transpose.extra_hints['direction'] = 'up'
        nhwc2nchw_transpose.extra_hints['direction'] = 'down'

        ops[0].inputs[input_idx] = new_input
        ops[-1].outputs[output_idx] = new_output

        return [nchw2nhwc_transpose] + ops + [nhwc2nchw_transpose]

    def wrap_ops_with_last_dim_transposes(
        self, ops: typing.List[tfl.BaseOperator], dim: int, input_idx: int = 0, output_idx: int = 0
    ) -> typing.List[tfl.BaseOperator]:
        orig_input = ops[0].inputs[input_idx]
        orig_output = ops[-1].outputs[output_idx]

        assert len(orig_input.shape) == len(orig_output.shape), "Numbers of dimensions mismatch"

        n_dim = len(orig_input.shape)
        if n_dim == dim:
            return ops

        last_dim_perm = np.array([i for i in range(n_dim) if i != dim] + [dim], dtype='int32')
        rev_last_dim_perm = np.argsort(last_dim_perm).astype('int32')

        last_dim_perm_tensor = self.create_attr_tensor(last_dim_perm)
        rev_last_dim_perm_tensor = self.create_attr_tensor(rev_last_dim_perm)

        new_input = self.create_transform_tensor(
            np.transpose(orig_input.tensor, last_dim_perm), quantization=orig_input.quantization
        )
        new_output = self.create_transform_tensor(
            np.transpose(orig_output.tensor, last_dim_perm), quantization=orig_output.quantization
        )

        last_dim_transpose = tfl.TransposeOperator([orig_input, last_dim_perm_tensor], [new_input])
        rev_last_dim_transpose = tfl.TransposeOperator([new_output, rev_last_dim_perm_tensor], [orig_output])

        last_dim_transpose.extra_hints['direction'] = 'up'
        rev_last_dim_transpose.extra_hints['direction'] = 'down'

        ops[0].inputs[input_idx] = new_input
        ops[-1].outputs[output_idx] = new_output

        return [last_dim_transpose] + ops + [rev_last_dim_transpose]

    def handle_padding(self, pad_h, pad_w, pad_op_index, ops, ceil_mode=False):
        fill_nan = False
        if ceil_mode:
            input_tensor = ops[0].inputs[0]
            kernel_size = [ops[1].filterHeight, ops[1].filterWidth]
            stride = [ops[1].strideH, ops[1].strideW]
            padding = [pad_h, pad_w]

            input_size = [input_tensor.shape[2], input_tensor.shape[3]]

            if not all((i + 2 * p - k) % s == 0 for i, p, k, s in zip(input_size, padding, kernel_size, stride)):
                assert type(ops[1]) is tfl.MaxPool2dOperator, 'ceil_mode=True for AvgPool not supported'
                fill_nan = True
                ceil_pad = get_pool_ceil_padding(input_tensor, kernel_size, stride, padding)
                ceil_pad = list(np.add(ceil_pad, padding))

        if pad_h + pad_w > 0:
            pad = [[0, 0], [pad_h, pad_h], [pad_w, pad_w], [0, 0]]
            pad_tensor = self.create_attr_tensor(np.array(pad, dtype='int32'))

            pad_input = ops[pad_op_index - 1].outputs[0]

            inputs = [pad_input, pad_tensor]
            if type(ops[1]) is tfl.MaxPool2dOperator:
                constant_tensor = self.get_minimum_constant(pad_input)
                inputs.append(constant_tensor)
                pad_array = np.pad(pad_input.tensor, pad, constant_values=constant_tensor.tensor[0])
            else:
                pad_array = np.pad(pad_input.tensor, pad)

            pad_out = self.create_transform_tensor(pad_array, quantization=pad_input.quantization)
            ops[pad_op_index].inputs[0] = pad_out

            if len(inputs) > 2:
                pad_op = tfl.Padv2Operator(inputs, [pad_out])
            else:
                pad_op = tfl.PadOperator(inputs, [pad_out])
            ops.insert(pad_op_index, pad_op)

        if fill_nan:
            fill_nan_index = pad_op_index + 1 if pad_h + pad_w > 0 else pad_op_index
            pad = [[0, 0], [0, ceil_pad[0]], [0, ceil_pad[1]], [0, 0]]
            pad_tensor = self.create_attr_tensor(np.array(pad, dtype='int32'))
            pad_input = ops[fill_nan_index - 1].outputs[0]
            constant_tensor = self.get_minimum_constant(pad_input)

            pad_array = np.pad(pad_input.tensor, pad, constant_values=constant_tensor.tensor[0])
            pad_out = self.create_transform_tensor(pad_array, quantization=pad_input.quantization)
            ops[fill_nan_index].inputs[0] = pad_out

            pad_op = tfl.Padv2Operator([pad_input, pad_tensor, constant_tensor], [pad_out])
            ops.insert(fill_nan_index, pad_op)

    def get_minimum_constant(self, ref_tensor):
        if ref_tensor.quantization is not None:
            if self.q_type == np.uint8:
                nan = 0
                constant_arr = tfl.FakeQuantTensor(
                    np.zeros(1, dtype=ref_tensor.dtype),
                    ref_tensor.quantization.scale,
                    ref_tensor.quantization.zero_point,
                )
            else:
                nan = -128
                constant_arr = tfl.FakeQuantTensor(
                    np.array([-128], dtype=ref_tensor.dtype),
                    ref_tensor.quantization.scale,
                    ref_tensor.quantization.zero_point,
                )
        else:
            nan = np.finfo(np.float32).min
            constant_arr = np.array([nan], dtype='float32')

        constant_tensor = self.create_attr_tensor(constant_arr)
        return constant_tensor

    def handle_reduce(self, converter_class, input_args, graph_converter, transpose_opt, *args, **kwargs):
        input_tensor = self.find_or_create_input(0, graph_converter)

        if 'dim' in input_args and 'keepdim' in input_args:
            dims, keep_dim = self.input_tensors[1:3]
            if type(dims) not in (list, tuple):
                dims = [dims]
            if len(dims) == 0:
                dims = list(range(input_tensor.tensor.ndim))
                self.output_tensors[0] = self.output_tensors[0].view(1)
        else:
            dims = list(range(input_tensor.tensor.ndim))
            keep_dim = False
            self.output_tensors[0] = self.output_tensors[0].view(1)

        for idx, dim in enumerate(dims):
            if dim < 0:
                dims[idx] += input_tensor.tensor.ndim

        ops = []
        transpose = False

        if transpose_opt:
            # For some ops the codepath is optimized for nhwc.
            # For example, for tfl.Mean, if it is a pooling 2d op, consider wrapping it with transposes
            if len(input_tensor.shape) == 4 and keep_dim in (1, True):
                if dims == [2, 3]:
                    dims = [1, 2]
                    transpose = True
                elif dims == [3, 2]:
                    dims = [2, 1]
                    transpose = True

        dim_tensor = self.create_attr_tensor(np.array(dims, dtype='int32'))

        inputs = [input_tensor, dim_tensor]
        outputs = self.to_tfl_tensors(self.output_names, self.output_tensors)

        if len(outputs) > 1:
            log.warning(
                'Reduce ops like `torch.min` have multiple outputs. However, only the first '
                'output will be preserved in our converter. If you need that tensor, please '
                'use the `torch.argmin` instead.'
            )
            outputs = outputs[:1]

        if (
            hasattr(converter_class, '__init__')
            and 'keepDims' in inspect.signature(converter_class.__init__).parameters
        ):
            ops.append(converter_class(inputs, outputs, keep_dim, *args, **kwargs))
        else:
            if keep_dim:
                output_tensor = outputs[0]
                transform = self.create_transform_tensor(np.squeeze(output_tensor.tensor, tuple(dims)))
                ops.append(converter_class(inputs, [transform], *args, **kwargs))

                shape_tensor = self.create_attr_tensor(np.array(output_tensor.shape, dtype='int32'))
                ops.append(tfl.ReshapeOperator([transform, shape_tensor], [output_tensor], shape_tensor.tensor))
            else:
                ops.append(converter_class(inputs, outputs, *args, **kwargs))

        if transpose:
            if keep_dim:
                ops = self.wrap_ops_with_nhwc_nchw_transposes(ops)
            else:
                orig_input = ops[0].inputs[0]

                nchw2nhwc_perm = np.array([0, 2, 3, 1], dtype='int32')
                nchw2nhwc_perm_tensor = self.create_attr_tensor(nchw2nhwc_perm)

                new_input = self.create_transform_tensor(
                    np.transpose(orig_input.tensor, nchw2nhwc_perm), quantization=orig_input.quantization
                )

                nchw2nhwc_transpose = tfl.TransposeOperator([orig_input, nchw2nhwc_perm_tensor], [new_input])

                ops[0].inputs[0] = new_input
                ops.insert(0, nchw2nhwc_transpose)

        for op in ops:
            graph_converter.add_operator(op)

    def quantize_scalar_tensor(self, tensor: torch.Tensor):
        assert tensor.numel() == 1
        assert tensor.dtype == torch.float32
        if not tensor.is_nonzero():
            if self.q_type in (np.uint8, np.int16):
                return torch.quantize_per_tensor(tensor, 0.5, 128, torch.quint8)
            elif self.q_type == np.int8:
                return torch.quantize_per_tensor(tensor, 0.5, 0, torch.qint8)
        elif (torch.sign(tensor) < 0).all():
            if self.q_type == np.uint8:
                return torch.quantize_per_tensor(tensor, -tensor[0] / 127, 255, torch.quint8)
            elif self.q_type == np.int8:
                return torch.quantize_per_tensor(tensor, -tensor[0] / 127, 0, torch.qint8)
            elif self.q_type == np.int16:
                return torch.quantize_per_tensor(tensor, -tensor[0] / 127, 128, torch.quint8)
        else:
            if self.q_type == np.uint8:
                return torch.quantize_per_tensor(tensor, tensor[0] / 127, 0, torch.quint8)
            elif self.q_type == np.int8:
                return torch.quantize_per_tensor(tensor, tensor[0] / 127, 0, torch.qint8)
            elif self.q_type == np.int16:
                return torch.quantize_per_tensor(tensor, tensor[0] / 127, 128, torch.quint8)

    def torch_tensor_from_scalar(self, ref_tensor: torch.Tensor, src_tensor: torch.Tensor):
        tgt_tensor = src_tensor
        if not isinstance(src_tensor, torch.Tensor):
            if ref_tensor.is_quantized:
                tgt_tensor = torch.quantize_per_tensor(
                    torch.tensor([src_tensor], dtype=torch.float32),
                    ref_tensor.q_scale(),
                    ref_tensor.q_zero_point(),
                    ref_tensor.dtype,
                )
            else:
                tgt_tensor = torch.tensor([src_tensor], dtype=ref_tensor.dtype)
        return tgt_tensor


def get_prop_from_node(node, prop, assert_type=None, return_type=False):
    output_name = next(node.outputs()).debugName()
    if prop in node.attributeNames():
        vk = node.kindOf(prop)
        if assert_type is not None and vk != assert_type:
            return None

        if vk == 'i':
            v = getattr(node, vk)(prop)
        elif vk == 'f':
            v = getattr(node, vk)(prop)
        elif vk == 's':
            v = getattr(node, vk)(prop)
        elif vk == 'g':
            v = getattr(node, vk)(prop)
        elif vk == 't':
            v = getattr(node, vk)(prop)
            if v.dtype == torch.float64:
                log.warning(
                    f'{output_name} is of type float64, which is unsupported in TFLite, trying to downcast to float32'
                )
                v = v.to(dtype=torch.float32)
        elif node.output().type().isSubtypeOf(torch._C.ListType.ofInts()) or node.output().type().isSubtypeOf(
            torch._C.ListType.ofFloats()
        ):
            v = node.output().toIValue()
        elif vk == 'ival':
            v = node.output().toIValue()
        else:
            log.warning(f'Skip unsupported constant generation for {output_name}, type: {vk}')
            raise StopIteration
    else:
        v = None
        vk = None

    if return_type:
        return v, vk
    else:
        return v


def fp32_to_bits(val):
    b = np.float32(val).tobytes()
    return np.frombuffer(b, dtype='uint32')[0]


def get_pool_ceil_padding(input, kernel_size, stride, padding):
    # Copied from the PyTorch repo
    # https://github.com/pytorch/pytorch/blob/master/torch/onnx/symbolic_opset9.py
    sizes = input.shape
    dim = sizes[-len(padding) :] if sizes is not None else None
    ceiled_output_dim = [
        int(math.ceil((dim[i] + 2 * padding[i] - kernel_size[i]) / float(stride[i]))) + 1
        for i in range(0, len(padding))
    ]
    # ensure last pooling starts inside
    ceiled_output_dim = [
        ceiled_output_dim[i] - 1
        if (((ceiled_output_dim[i] - 1) * stride[i]) >= (dim[i] + padding[i]))
        else ceiled_output_dim[i]
        for i in range(0, len(ceiled_output_dim))
    ]
    padding_ceil = [
        0
        if (stride[i] == 1)
        else (kernel_size[i] - (dim[i] + 2 * padding[i] - ((ceiled_output_dim[i] - 1) * stride[i] + 1)))
        for i in range(0, len(padding))
    ]
    # ensure padding is not > kernel_size
    padding_ceil = [
        (int(padding_ceil[i]) if padding_ceil[i] < kernel_size[i] - 1 else int(kernel_size[i] - 1))
        if ((padding_ceil[i] + 2 * padding[i]) >= (kernel_size[i]))
        else int(padding_ceil[i])
        for i in range(0, len(padding_ceil))
    ]
    return padding_ceil


class NoTrackOperator(OperatorConverter):
    def parse(self, node, attrs, args, graph_converter):
        super().parse(node, attrs, args, graph_converter)

        self.run(node)


class TrackQParamsOperator(OperatorConverter):
    def parse(self, node, attrs, args, graph_converter):
        super().parse(node, attrs, args, graph_converter)

        self.run(node)

        t = self.find_or_create_input(0, graph_converter)
        graph_converter.q_mapping[self.output_names[0]] = t


class TrackRevQParamsOperator(OperatorConverter):
    def parse(self, node, attrs, args, graph_converter):
        super().parse(node, attrs, args, graph_converter)

        self.run(node)

        t = self.to_tfl_tensors(self.output_names, self.output_tensors)[0]
        graph_converter.rev_q_mapping[self.input_names[0]] = t


class TrackConstantOperator(OperatorConverter):
    def parse(self, node, attrs, args, graph_converter):
        super().parse(node, attrs, args, graph_converter)

        self.run(node)

        t = self.find_or_create_input(0, graph_converter)
        graph_converter.constant_mapping[self.output_names[0]] = t


class PrimOperatorConverter(OperatorConverter):
    # prim::* ops needs custom implementation
    def run(self, node):
        pass