transformable.py

from abc import abstractmethod

from .base import BaseOperator, QuantizationParameters, Tensor
from .custom import MTKTransposeConvOperator
from . import generated_ops as tfl_ops

from ..base import ExtendedOperator
from ...schemas.tflite import schema_generated as tflite

import typing
import torch
import warnings

import numpy as np


class TransformableOperator(BaseOperator):
    def __init__(self, op: int, inputs: typing.List['Tensor'], outputs: typing.List['Tensor'], op_version: int):
        super().__init__(op, inputs, outputs, op_version=op_version)
        self.attr_count = 0
        self.transform_count = 0

    @abstractmethod
    def transform(self):
        pass

    def create_attr_tensor(self, tensor, name=None, quantization=None):
        if name is None:
            if self.attr_count == 0:
                name = self.outputs[0].name + '_te_attr'
            else:
                name = self.outputs[0].name + f'_te_attr_{self.attr_count}'
            self.attr_count += 1
        return Tensor(tensor, name, has_buffer=True, quantization=quantization)

    def create_transform_tensor(self, tensor, name=None, quantization=None):
        if name is None:
            if self.transform_count == 0:
                name = self.outputs[0].name + '_te_transform'
            else:
                name = self.outputs[0].name + f'_te_transform_{self.transform_count}'
            self.transform_count += 1
        return Tensor(tensor, name, has_buffer=False, quantization=quantization)

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

        if orig_input.tensor.ndim == 4:
            nhwc2nchw_perm = np.array([0, 3, 1, 2], dtype='int32')
            nchw2nhwc_perm = np.array([0, 2, 3, 1], dtype='int32')
        elif orig_input.tensor.ndim == 5:
            nhwc2nchw_perm = np.array([0, 4, 1, 2, 3], dtype='int32')
            nchw2nhwc_perm = np.array([0, 2, 3, 4, 1], dtype='int32')
        else:
            assert False, f'Don\'t know how to wrap tranposes for {orig_input.tensor.ndim}d tensors'

        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_ops.TransposeOperator([orig_input, nchw2nhwc_perm_tensor], [new_input])
        nhwc2nchw_transpose = tfl_ops.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]


class BatchNormOperator(TransformableOperator):
    input_index = 0
    weight_index = 1
    bias_index = 2
    running_mean_index = 3
    running_variance_index = 4

    output_index = 0

    def __init__(
        self,
        inputs: typing.List['Tensor'],
        outputs: typing.List['Tensor'],
        eps: float,
        quantization: typing.Optional[QuantizationParameters] = None,
        fusedActivationFunction=tflite.ActivationFunctionType.NONE,
    ):
        super().__init__(ExtendedOperator.BATCH_NORM, inputs, outputs, 1)
        self.eps = eps
        self.fusedActivationFunction = fusedActivationFunction

    def transform(self, graph_converter, mapping):
        assert all((x.buffer is not None for x in self.inputs[1:]))

        w, b, mean, var = [
            self.inputs[i]
            for i in (self.weight_index, self.bias_index, self.running_mean_index, self.running_variance_index)
        ]

        inv = 1 / np.sqrt(var.tensor + self.eps)
        new_w = inv * w.tensor
        new_b = b.tensor - mean.tensor * new_w

        inp = self.inputs[0]

        new_shape = [1] + [new_w.shape[0]] + [1] * (inp.tensor.ndim - 2)

        new_w = new_w.reshape(new_shape)
        new_b = new_b.reshape(new_shape)

        weight = self.create_attr_tensor(new_w)
        bias = self.create_attr_tensor(new_b)

        new_inp = inp
        if inp.quantization is not None:
            new_inp = self.create_transform_tensor(inp.tensor)
            graph_converter.add_operator(tfl_ops.DequantizeOperator([inp], [new_inp]))

        mul_out = self.create_transform_tensor(new_inp.tensor * weight.tensor)
        graph_converter.add_operator(tfl_ops.MulOperator([new_inp, weight], [mul_out]))

        if inp.quantization is not None:
            add_out = self.create_transform_tensor(mul_out.tensor + bias.tensor)
        else:
            add_out = self.outputs[self.output_index]

        graph_converter.add_operator(
            tfl_ops.AddOperator([mul_out, bias], [add_out], fusedActivationFunction=self.fusedActivationFunction),
            transform=True,
        )

        if inp.quantization is not None:
            quant_out = self.outputs[self.output_index]
            graph_converter.add_operator(tfl_ops.QuantizeOperator([add_out], [quant_out]), transform=True)

        graph_converter.try_restore_edges(mapping)


class GenericConvOperator(TransformableOperator):
    input_index = 0
    weight_index = 1
    bias_index = 2

    output_index = 0

    stride: typing.List[int]
    padding: typing.List[int]
    dilation: typing.List[int]
    transpose: bool
    output_padding: typing.List[int]
    groups: int

    fusedActivationFunction: tflite.ActivationFunctionType

    def __init__(
        self,
        inputs: typing.List['Tensor'],
        outputs: typing.List['Tensor'],
        stride: typing.List[int],
        padding: typing.List[int],
        dialation: typing.List[int],
        output_padding: typing.List[int],
        groups: int,
        fusedActivationFunction=tflite.ActivationFunctionType.NONE,
    ):
        super().__init__(ExtendedOperator.GENERIC_CONV, inputs, outputs, 1)
        self.stride = stride
        self.padding = padding
        self.dilation = dialation
        self.output_padding = output_padding
        self.groups = groups

        self.fusedActivationFunction = fusedActivationFunction

    def transform(self, graph_converter, mapping):
        input_tensor = self.inputs[0]
        weight_tensor = self.inputs[1]

        input_dim = len(input_tensor.shape)
        weight_dim = len(weight_tensor.shape)

        prev_ops = []
        next_ops = []

        if weight_dim == 3 or input_dim == 3:
            reshape_input_size = 1
            reshape_output_size = 1
            if weight_dim == 3:
                self.stride.insert(0, 1)
                self.padding.insert(0, 0)
                self.dilation.insert(0, 1)
                self.output_padding.insert(0, 0)
                reshape_input_size = 2

            reshape_outputs = [
                self.create_transform_tensor(
                    np.expand_dims(t.tensor, 2),
                    name=f'{self.outputs[0].name}_{t.name}_4d_input',
                    quantization=t.quantization,
                )
                for t in self.inputs[:reshape_input_size]
            ]
            reshape_attrs = [self.create_attr_tensor(np.array(t.shape, dtype='int32')) for t in reshape_outputs]
            reshape_ops = [
                tfl_ops.ReshapeOperator([old, attr], [new], attr.tensor)
                for old, new, attr in zip(self.inputs[:reshape_input_size], reshape_outputs, reshape_attrs)
            ]

            for op in reshape_ops:
                op.extra_hints['direction'] = 'up'

            prev_ops.extend(reshape_ops)

            conv_outputs = [
                self.create_transform_tensor(
                    np.expand_dims(self.outputs[i].tensor, 2),
                    name=f'{self.outputs[i].name}_4d_output',
                    quantization=self.outputs[i].quantization,
                )
                for i in range(reshape_output_size)
            ]
            conv_attrs = [
                self.create_attr_tensor(np.array(t.shape, dtype='int32')) for t in self.outputs[:reshape_output_size]
            ]
            conv_ops = [
                tfl_ops.ReshapeOperator([old, attr], [new], attr.tensor)
                for old, new, attr in zip(conv_outputs, self.outputs[:reshape_output_size], conv_attrs)
            ]

            for op in conv_ops:
                op.extra_hints['direction'] = 'down'

            next_ops.extend(conv_ops)

            self.inputs = reshape_outputs + self.inputs[reshape_input_size:]
            self.outputs = conv_outputs + self.outputs[reshape_output_size:]

            weight_tensor = self.inputs[1]
        elif weight_dim not in (4, 5):
            assert False, "Only Conv[Transpose]1d/2d/3d is supported"

        if weight_tensor.shape[1] == 1 and weight_tensor.shape[0] == self.groups:
            if weight_dim in (3, 4):
                conv_op = tfl_ops.DepthwiseConv2dOperator(
                    self.inputs,
                    self.outputs,
                    strideH=self.stride[0],
                    strideW=self.stride[1],
                    depthMultiplier=1,
                    dilationHFactor=self.dilation[0],
                    dilationWFactor=self.dilation[1],
                    fusedActivationFunction=self.fusedActivationFunction,
                    padding=tflite.Padding.VALID,
                )
            else:
                assert False, "Only DepthwiseConv1d/2d is supported"
        else:
            if input_tensor.shape[1] != weight_tensor.shape[1]:
                warnings.warn(
                    'Group conv is not supported if official tflite interpreter is used. If that is the case for you,'
                    ' plese pass in `group_conv_rewrite=True`. If you want to run the model with TFLite micro, then you'
                    ' may also need to pass in `tflite_micro_rewrite=True`'
                )
            if weight_dim in (3, 4):
                conv_op = tfl_ops.Conv2dOperator(
                    self.inputs,
                    self.outputs,
                    strideH=self.stride[0],
                    strideW=self.stride[1],
                    dilationHFactor=self.dilation[0],
                    dilationWFactor=self.dilation[1],
                    fusedActivationFunction=self.fusedActivationFunction,
                    padding=tflite.Padding.VALID,
                )
            else:
                conv_op = tfl_ops.Conv3dOperator(
                    self.inputs,
                    self.outputs,
                    strideD=self.stride[0],
                    strideH=self.stride[1],
                    strideW=self.stride[2],
                    dilationDFactor=self.dilation[0],
                    dilationHFactor=self.dilation[1],
                    dilationWFactor=self.dilation[2],
                    fusedActivationFunction=self.fusedActivationFunction,
                    padding=tflite.Padding.VALID,
                )

        ops = self.wrap_ops_with_nhwc_nchw_transposes([conv_op])
        conv_op = ops[1]

        # Pad handling
        if sum(self.padding) > 0:
            if weight_dim in (3, 4):
                pad_h = self.padding[0]
                pad_w = self.padding[1]

                pad = [[0, 0], [pad_h, pad_h], [pad_w, pad_w], [0, 0]]
            else:
                pad_d = self.padding[0]
                pad_h = self.padding[1]
                pad_w = self.padding[2]

                pad = [[0, 0], [pad_d, pad_d], [pad_h, pad_h], [pad_w, pad_w], [0, 0]]

            pad_tensor = self.create_attr_tensor(np.array(pad, dtype='int32'))

            pad_input = ops[0].outputs[0]
            pad_array = np.pad(pad_input.tensor, pad)
            pad_out = self.create_transform_tensor(pad_array, quantization=pad_input.quantization)
            ops[1].inputs[0] = pad_out

            pad_op = tfl_ops.PadOperator([pad_input, pad_tensor], [pad_out])
            ops.insert(1, pad_op)

        # Weight handling
        weight = conv_op.inputs[1]
        if conv_op.op.code == tflite.BuiltinOperator.DEPTHWISE_CONV_2D:
            nchw2chwn_perm = np.array([1, 2, 3, 0], dtype='int32')
            nchw2chwn_perm_tensor = self.create_attr_tensor(nchw2chwn_perm)
            weight_q = weight.quantization
            if weight_q is not None and weight_q.dim is not None:
                new_dim = np.nonzero(nchw2chwn_perm == weight_q.dim)[0][0]
                weight_q = QuantizationParameters(weight_q.scale, weight_q.zero_point, new_dim)
            reordered_weight = self.create_transform_tensor(
                np.transpose(weight.tensor, nchw2chwn_perm), quantization=weight_q
            )
            conv_op.inputs[1] = reordered_weight
            reorder_op = tfl_ops.TransposeOperator([weight, nchw2chwn_perm_tensor], [reordered_weight])
        else:
            if weight_dim in (3, 4):
                nchw2nhwc_perm = np.array([0, 2, 3, 1], dtype='int32')
                nchw2nhwc_perm_tensor = self.create_attr_tensor(nchw2nhwc_perm)
            else:
                nchw2nhwc_perm = np.array([2, 3, 4, 1, 0], dtype='int32')
                nchw2nhwc_perm_tensor = self.create_attr_tensor(nchw2nhwc_perm)
            weight_q = weight.quantization
            if weight_q is not None and weight_q.dim is not None:
                new_dim = np.nonzero(nchw2nhwc_perm == weight_q.dim)[0][0]
                weight_q = QuantizationParameters(weight_q.scale, weight_q.zero_point, new_dim)
            reordered_weight = self.create_transform_tensor(
                np.transpose(weight.tensor, nchw2nhwc_perm), quantization=weight_q
            )
            conv_op.inputs[1] = reordered_weight
            reorder_op = tfl_ops.TransposeOperator([weight, nchw2nhwc_perm_tensor], [reordered_weight])
        ops.insert(1, reorder_op)

        # Bias handling
        kernel_num = self.inputs[1].shape[0]
        if conv_op.op.code in (tflite.BuiltinOperator.DEPTHWISE_CONV_2D, tflite.BuiltinOperator.CONV_3D):
            kernel_num = self.inputs[1].shape[-1]

        if len(conv_op.inputs) == 2 or conv_op.inputs[2] is None:
            if conv_op.inputs[0].dtype == np.dtype('float32'):
                bias = np.zeros((kernel_num,), dtype='float32')
                q_args = None
            else:
                bias = np.zeros((kernel_num,), dtype='int32')

                per_tensor = weight_tensor.quantization.dim is None

                # Bias handling
                if per_tensor:
                    bias_scale = input_tensor.quantization.scale * weight_tensor.quantization.scale
                    bias_zero_point = 0
                    bias_dim = None
                else:
                    bias_scale = [input_tensor.quantization.scale * s for s in weight_tensor.quantization.scale]
                    bias_zero_point = [0] * len(bias_scale)
                    bias_dim = 0

                q_args = QuantizationParameters(bias_scale, bias_zero_point, bias_dim)

            conv_op.inputs.append(self.create_attr_tensor(bias, quantization=q_args))
        elif conv_op.inputs[2].shape[0] != kernel_num and conv_op.inputs[2].shape[0] == 1:
            if conv_op.inputs[0].dtype == np.float32:
                bias = torch.tensor([conv_op.inputs[2][0]] * kernel_num, dtype='float32')
            else:
                bias = torch.tensor([conv_op.inputs[2][0]] * kernel_num, dtype='int32')

            conv_op.inputs[2] = self.create_attr_tensor(bias)

        ops = prev_ops + ops + next_ops

        for op in ops:
            graph_converter.add_operator(op, transform=True)

        graph_converter.try_restore_edges(mapping)

        for op in ops[:-1]:
            output_name = op.outputs[0].name
            node_name = graph_converter.tensor_node_map[output_name]
            node = graph_converter.graph.vs.find(name=node_name)
            assert node.outdegree() > 0, (
                'The following node should be a part of the transformable node,                 but the outdegree of'
                f' it is zero. {node}'
            )
            next_node = graph_converter.graph.vs[node.out_edges()[0].target]
            assert next_node['node_type'] != ExtendedOperator.CONSTANT_NODE


class GenericTransposeConvOperator(TransformableOperator):
    input_index = 0
    weight_index = 1
    bias_index = 2

    output_index = 0

    stride: typing.List[int]
    padding: typing.List[int]
    dilation: typing.List[int]
    transpose: bool
    output_padding: typing.List[int]
    groups: int

    enable_mtk_ops: bool
    conv_transpose_with_bias: bool

    fusedActivationFunction: tflite.ActivationFunctionType

    def __init__(
        self,
        inputs: typing.List['Tensor'],
        outputs: typing.List['Tensor'],
        stride: typing.List[int],
        padding: typing.List[int],
        dilation: typing.List[int],
        output_padding: typing.List[int],
        groups: int,
        enable_mtk_ops: bool = False,
        conv_transpose_with_bias: bool = True,
        fusedActivationFunction=tflite.ActivationFunctionType.NONE,
    ):
        super().__init__(ExtendedOperator.GENERIC_DECONV, inputs, outputs, 1)
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.output_padding = output_padding
        self.groups = groups
        self.enable_mtk_ops = enable_mtk_ops
        self.conv_transpose_with_bias = conv_transpose_with_bias

        self.fusedActivationFunction = fusedActivationFunction

    def transform(self, graph_converter, mapping):
        input_tensor = self.inputs[0]
        weight_tensor = self.inputs[1]
        output_tensor = self.outputs[0]

        input_dim = len(input_tensor.shape)
        weight_dim = len(weight_tensor.shape)

        prev_ops = []
        next_ops = []

        if weight_dim == 3 or input_dim == 3:
            self.stride.insert(0, 1)
            self.padding.insert(0, 0)
            self.dilation.insert(0, 1)
            self.output_padding.insert(0, 0)

            reshape_outputs = [
                self.create_transform_tensor(
                    np.expand_dims(t.tensor, 2),
                    name=f'{self.outputs[0].name}_{t.name}_4d_input',
                    quantization=t.quantization,
                )
                for t in self.inputs[:2]
            ]
            reshape_attrs = [self.create_attr_tensor(np.array(t.shape, dtype='int32')) for t in reshape_outputs]
            reshape_ops = [
                tfl_ops.ReshapeOperator([old, attr], [new], attr.tensor)
                for old, new, attr in zip(self.inputs[:2], reshape_outputs, reshape_attrs)
            ]

            for op in reshape_ops:
                op.extra_hints['direction'] = 'up'

            if weight_dim == 3 and input_dim == 3:
                prev_ops.extend(reshape_ops)
            elif weight_dim == 3:
                prev_ops.append(reshape_ops[1])
            else:
                prev_ops.append(reshape_ops[0])

            conv_outputs = [
                self.create_transform_tensor(
                    np.expand_dims(self.outputs[0].tensor, 2),
                    name=f'{self.outputs[0].name}_4d_output',
                    quantization=self.outputs[0].quantization,
                )
            ]
            conv_attrs = [self.create_attr_tensor(np.array(t.shape, dtype='int32')) for t in self.outputs[:1]]
            conv_ops = [
                tfl_ops.ReshapeOperator([old, attr], [new], attr.tensor)
                for old, new, attr in zip(conv_outputs, self.outputs[:1], conv_attrs)
            ]

            for op in conv_ops:
                op.extra_hints['direction'] = 'down'

            next_ops.extend(conv_ops)

            if weight_dim == 3 and input_dim == 3:
                self.inputs = reshape_outputs + self.inputs[2:]
            elif weight_dim == 3:
                self.inputs = self.inputs[0:1] + reshape_outputs[1:2] + self.inputs[1:]
            else:
                self.inputs = reshape_outputs[0:1] + self.inputs[1:]
            self.outputs = conv_outputs + self.outputs[1:]

            weight_tensor = self.inputs[1]
        elif weight_dim not in (4, 5):
            assert False, "Only Conv[Transpose]1d/2d/3d is supported"

        if output_tensor.shape[1] != weight_tensor.shape[1]:
            warnings.warn(
                'Group transposed conv is not supported if official tflite interpreter is used. If that is the case'
                ' for you, plese pass in `group_conv_rewrite=True`. If you want to run the model with TFLite micro,'
                ' then you may also need to pass in `tflite_micro_rewrite=True`'
            )

        if weight_dim in (3, 4):
            assert all((x == 1 for x in self.dilation)), "Only dilation=1 is supported for conv_transpose2d"
            if self.enable_mtk_ops:
                conv_op = MTKTransposeConvOperator(
                    self.inputs[:2][::-1],
                    self.outputs,
                    depth_multiplier=1,
                    dilation_height_factor=self.dilation[0],
                    dilation_width_factor=self.dilation[1],
                    padding_type=tflite.Padding.VALID,
                    stride_height=self.stride[0],
                    stride_width=self.stride[1],
                )
            else:
                conv_op = tfl_ops.TransposeConvOperator(
                    self.inputs[:2][::-1],
                    self.outputs,
                    strideH=self.stride[0],
                    strideW=self.stride[1],
                    padding=tflite.Padding.VALID,
                    fusedActivationFunction=self.fusedActivationFunction,
                )
        else:
            conv_op = tfl_ops.Conv3dTransposeOperator(
                self.inputs[:2][::-1],
                self.outputs,
                strideD=self.stride[0],
                strideH=self.stride[1],
                strideW=self.stride[2],
                dilationDFactor=self.dilation[0],
                dilationHFactor=self.dilation[1],
                dilationWFactor=self.dilation[2],
                padding=tflite.Padding.VALID,
                fusedActivationFunction=self.fusedActivationFunction,
            )

        ops = self.wrap_ops_with_nhwc_nchw_transposes([conv_op], input_idx=1)

        # Pad handling
        output_shape = conv_op.outputs[0].shape
        if sum(self.padding) > 0:
            if weight_dim in (3, 4):
                pad_h = self.padding[0]
                pad_w = self.padding[1]

                start = np.array([0, pad_h, pad_w, 0], dtype='int32')

                pad_sizes = ((0, 0), (pad_h, pad_h), (pad_w, pad_w), (0, 0))
            else:
                pad_d = self.padding[0]
                pad_h = self.padding[1]
                pad_w = self.padding[2]

                start = np.array([0, pad_d, pad_h, pad_w, 0], dtype='int32')

                pad_sizes = ((0, 0), (pad_d, pad_d), (pad_h, pad_h), (pad_w, pad_w), (0, 0))

            size = np.array(ops[1].outputs[0].shape, dtype='int32')

            start_tensor = self.create_attr_tensor(start)
            size_tensor = self.create_attr_tensor(size)

            slice_out = ops[1].outputs[0]
            pad_array = np.pad(self.outputs[0].tensor, pad_sizes)
            slice_input = self.create_transform_tensor(pad_array, quantization=self.outputs[0].quantization)
            ops[1].outputs[0] = slice_input

            slice_op = tfl_ops.SliceOperator([slice_input, start_tensor, size_tensor], [slice_out])
            output_shape = slice_input.shape
            ops.insert(2, slice_op)

        # Output shape handling
        output_shape_tensor = self.create_attr_tensor(np.array(output_shape, dtype='int32'))
        conv_op.inputs.insert(0, output_shape_tensor)

        # Weight handling
        weight = conv_op.inputs[1]
        if weight_dim in (3, 4):
            nchw2chwn_perm = np.array([1, 2, 3, 0], dtype='int32')
        else:
            nchw2chwn_perm = np.array([2, 3, 4, 1, 0], dtype='int32')
        nchw2chwn_perm_tensor = self.create_attr_tensor(nchw2chwn_perm)
        reordered_weight = self.create_transform_tensor(
            np.transpose(weight.tensor, nchw2chwn_perm), quantization=weight.quantization
        )
        conv_op.inputs[1] = reordered_weight
        reorder_op = tfl_ops.TransposeOperator([weight, nchw2chwn_perm_tensor], [reordered_weight])
        ops.insert(1, reorder_op)

        # Bias handling
        if self.enable_mtk_ops or self.conv_transpose_with_bias:
            kernel_num = output_tensor.shape[1]

            if len(self.inputs) > 2 and self.inputs[2].shape[0] != kernel_num and self.inputs[2].shape[0] == 1:
                if conv_op.inputs[-1].dtype == np.float32:
                    bias = torch.tensor([self.inputs[2][0]] * kernel_num, dtype='float32')
                else:
                    bias = torch.tensor([self.inputs[2][0]] * kernel_num, dtype='int32')

                conv_op.inputs.append(self.create_attr_tensor(bias))

            else:
                if len(self.inputs) == 2 or self.inputs[2] is None:
                    if conv_op.inputs[-1].dtype == np.dtype('float32'):
                        bias = np.zeros((kernel_num,), dtype='float32')
                        q_args = None
                    else:
                        bias = np.zeros((kernel_num,), dtype='int32')
                else:
                    bias = self.inputs[2].tensor

                q_args = None
                if bias.dtype != np.dtype('float32'):
                    per_tensor = weight_tensor.quantization.dim is None

                    # Bias handling
                    if per_tensor:
                        bias_scale = input_tensor.quantization.scale * weight_tensor.quantization.scale
                        bias_zero_point = 0
                        bias_dim = None
                    else:
                        bias_scale = [input_tensor.quantization.scale * s for s in weight_tensor.quantization.scale]
                        bias_zero_point = [0] * len(bias_scale)
                        bias_dim = 0

                    q_args = QuantizationParameters(bias_scale, bias_zero_point, bias_dim)

                conv_op.inputs.append(self.create_attr_tensor(bias, quantization=q_args))
        else:
            if len(self.inputs) > 2 and self.inputs[2] is not None:
                bias_tensor = self.inputs[2]
                add_out = ops[-2].outputs[0]
                bias_transform = self.create_transform_tensor(
                    add_out.tensor.copy(), quantization=self.outputs[0].quantization
                )
                ops[-2].outputs[0] = bias_transform
                ops.insert(len(ops) - 1, tfl_ops.AddOperator([bias_transform, bias_tensor], [add_out]))

        ops = prev_ops + ops + next_ops

        for op in ops:
            graph_converter.add_operator(op)

        graph_converter.try_restore_edges(mapping)