# Copyright (C) 2023-2024, Advanced Micro Devices, Inc.
# 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 FINN 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 AND CONTRIBUTORS "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 HOLDER 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.

"""Transformations to map ONNX operators to FINN hardware layers."""

import numpy as np
import qonnx.core.data_layout as DataLayout
from onnx import NodeProto, TensorProto, helper
from qonnx.core.datatype import DataType
from qonnx.core.modelwrapper import ModelWrapper
from qonnx.custom_op.registry import getCustomOp
from qonnx.transformation.base import Transformation
from qonnx.transformation.general import SortGraph
from qonnx.transformation.infer_datatypes import InferDataTypes
from qonnx.transformation.infer_shapes import InferShapes
from qonnx.util.basic import get_by_name, remove_by_name
from qonnx.util.onnx import nchw_to_nhwc

# Module containing specializations of elementwise binary operations
import finn.custom_op.fpgadataflow.elementwise_binary as elementwise_binary

# Base class for all FINN custom ops, here just used for type-hinting
from finn.custom_op.fpgadataflow.hwcustomop import HWCustomOp
from finn.util.logging import log


class InferConvInpGen(Transformation):
    """Convert Im2Col layers to ConvolutionInputGenerator layers."""

    def __init__(self):
        """Initialize the transformation."""
        super().__init__()

    def apply(self, model):
        """Apply the transformation to infer ConvolutionInputGenerator layers."""
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for n in graph.node:
            node_ind += 1
            if n.op_type == "Im2Col":
                i2c_input = n.input[0]
                i2c_output = n.output[0]
                i2c_in_shape = model.get_tensor_shape(i2c_input)
                i2c_out_shape = model.get_tensor_shape(i2c_output)
                dt = model.get_tensor_datatype(i2c_input)
                if not dt.is_integer():
                    log.warning(f"{n.name} : Input is not int. Can't infer ConvInpGen.")
                    continue
                i2c_inst = getCustomOp(n)
                stride_h, stride_w = i2c_inst.get_nodeattr("stride")
                k_h, k_w = i2c_inst.get_nodeattr("kernel_size")
                pad_attr = i2c_inst.get_nodeattr("pad_amount")
                pad_h = pad_attr[0] + pad_attr[2]
                pad_w = pad_attr[1] + pad_attr[3]
                dilation_h, dilation_w = i2c_inst.get_nodeattr("dilations")
                pad_val = i2c_inst.get_nodeattr("pad_value")
                depthwise = i2c_inst.get_nodeattr("depthwise")
                ifm_ch = i2c_in_shape[-1]
                ifm_dim_h = i2c_in_shape[1]
                ifm_dim_w = i2c_in_shape[2]
                ofm_dim_h = i2c_out_shape[1]
                ofm_dim_w = i2c_out_shape[2]

                # default params for ConvolutionInputGenerator
                ConvInpGen_node_idx = node_ind
                ConvInpGen_input = i2c_input
                ConvInpGen_idim_h = ifm_dim_h
                ConvInpGen_idim_w = ifm_dim_w

                if pad_h > 0 or pad_w > 0:
                    assert pad_val == 0, (
                        "%s : FMPadding_Batch doesn't currently support pad_val!= 0" % n.name
                    )

                    odim_padding_h = ifm_dim_h + pad_h
                    odim_padding_w = ifm_dim_w + pad_w

                    padding_out = helper.make_tensor_value_info(
                        model.make_new_valueinfo_name(),
                        TensorProto.FLOAT,
                        (1, odim_padding_h, odim_padding_w, ifm_ch),
                    )
                    graph.value_info.append(padding_out)
                    padding_out = padding_out.name
                    model.set_tensor_datatype(padding_out, dt)

                    ConvInpGen_node_idx += 1
                    ConvInpGen_input = padding_out
                    ConvInpGen_idim_h = odim_padding_h
                    ConvInpGen_idim_w = odim_padding_w

                    padding_node = helper.make_node(
                        "FMPadding",
                        [i2c_input],
                        [padding_out],
                        domain="finn.custom_op.fpgadataflow",
                        backend="fpgadataflow",
                        ImgDim=[ifm_dim_h, ifm_dim_w],
                        Padding=pad_attr,
                        NumChannels=ifm_ch,
                        inputDataType=dt.name,
                        SIMD=ifm_ch,
                        name="FMPadding_Batch_" + n.name,
                    )
                    graph.node.insert(node_ind, padding_node)

                is_1D = (ifm_dim_h == 1) or (ifm_dim_w == 1)
                ConvInpGen_node = helper.make_node(
                    "ConvolutionInputGenerator",
                    [ConvInpGen_input],
                    [i2c_output],
                    domain="finn.custom_op.fpgadataflow",
                    backend="fpgadataflow",
                    ConvKernelDim=[k_h, k_w],
                    IFMChannels=ifm_ch,
                    IFMDim=[ConvInpGen_idim_h, ConvInpGen_idim_w],
                    OFMDim=[ofm_dim_h, ofm_dim_w],
                    SIMD=ifm_ch,
                    Stride=[stride_h, stride_w],
                    Dilation=[dilation_h, dilation_w],
                    inputDataType=dt.name,
                    outputDataType=dt.name,
                    depthwise=depthwise,
                    is1D=is_1D,
                    name="ConvolutionInputGenerator_" + n.name,
                )
                graph.node.insert(ConvInpGen_node_idx, ConvInpGen_node)
                # remove old nodes
                graph.node.remove(n)
                graph_modified = True
        if graph_modified:
            model = model.transform(InferShapes())
            model = model.transform(InferDataTypes())
        return (model, graph_modified)


class InferThresholdingLayer(Transformation):
    """Convert any MultiThreshold into a standalone thresholding layer."""

    def __init__(self):
        """Initialize the transformation."""
        super().__init__()

    def apply(self, model):
        """Apply the transformation to infer standalone thresholding layers."""
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for node in graph.node:
            node_ind += 1
            if node.op_type == "MultiThreshold":
                thl_input = node.input[0]
                thl_threshold = node.input[1]
                thl_output = node.output[0]
                thl_in_shape = model.get_tensor_shape(thl_input)
                thl_thres_shape = model.get_tensor_shape(thl_threshold)
                idt = model.get_tensor_datatype(thl_input)
                tdt = model.get_tensor_datatype(thl_threshold)

                # only infer layers where input and thresholds are integers or fp32
                idt_int = idt.is_integer()
                tdt_int = tdt.is_integer()
                idt_fp32 = idt == "FLOAT32"
                tdt_fp32 = tdt == "FLOAT32"
                if not (idt_int or idt_fp32):
                    continue
                if not (tdt_int or tdt_fp32):
                    continue

                # TODO: this should be removed in favor of proper layout handling in the frontend
                #  this workaround is currently still needed to handle standalone NCHW MTs at the
                #  input of the graph, e.g., for cnv (bnn-pynq) models
                # NOTE: careful, this requires proper layout annotation for input AND output,
                #  otherwise the graph could break because only one of two Transpose gets inserted

                # check layout of inputs/outputs, and convert if needed
                # check layout and convert if necessary
                thl_in_layout = model.get_tensor_layout(thl_input)
                if thl_in_layout == DataLayout.NCHW:
                    thl_input = nchw_to_nhwc(thl_input, model, node_ind)
                    node_ind += 1
                    thl_in_shape = model.get_tensor_shape(thl_input)

                # keep track of where we need to insert the HLS Op
                # it has to be ahead of the output transform
                insert_point = node_ind
                thl_output_layout = model.get_tensor_layout(thl_output)
                if thl_output_layout == DataLayout.NCHW:
                    thl_output = nchw_to_nhwc(thl_output, model, node_ind, reverse=True)
                    node_ind += 1

                # now safe to assume number of channels is in last dimension
                ifc = int(thl_in_shape[-1])
                # create node with no parallelization first
                pe = 1

                odt = model.get_tensor_datatype(thl_output)
                scale = getCustomOp(node).get_nodeattr("out_scale")
                assert scale == 1.0, (
                    node.name + ": MultiThreshold out_scale must be 1 for HLS conversion."
                )
                actval = getCustomOp(node).get_nodeattr("out_bias")
                assert int(actval) == actval, (
                    node.name + ": MultiThreshold out_bias must be integer for HLS conversion."
                )
                actval = int(actval)

                # a signed activation should always have a negative bias,
                # but BIPOLAR uses the -1 as 0 encoding so the assert does not apply
                if odt != DataType["BIPOLAR"]:
                    assert (not odt.signed()) or (actval < 0), (
                        node.name + ": Signed output requires actval < 0"
                    )

                new_node = helper.make_node(
                    "Thresholding",
                    [thl_input, thl_threshold],
                    [thl_output],
                    domain="finn.custom_op.fpgadataflow",
                    backend="fpgadataflow",
                    NumChannels=ifc,
                    PE=pe,
                    numSteps=thl_thres_shape[1],
                    inputDataType=idt.name,
                    weightDataType=tdt.name,
                    outputDataType=odt.name,
                    numInputVectors=list(thl_in_shape[:-1]),
                    ActVal=actval,
                    name="Thresholding_" + node.name,
                )

                graph.node.insert(insert_point, new_node)
                # remove old node
                graph.node.remove(node)
                graph_modified = True

        return (model, graph_modified)


class InferUpsample(Transformation):
    """Convert Upsample and Resize nodes to UpsampleNearestNeighbour nodes."""

    def apply(self, model):
        """Apply the transformation to infer UpsampleNearestNeighbour nodes."""
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for n in graph.node:
            node_ind += 1
            if n.op_type == "Upsample" or n.op_type == "Resize":
                # Extract mode and scales and input shape
                mode = get_by_name(n.attribute, "mode").s.decode("ascii")
                if n.op_type == "Upsample":
                    scales = model.get_initializer(n.input[1])
                else:
                    if len(n.input) == 2:
                        # Resize version 10
                        scales = model.get_initializer(n.input[1])
                    elif len(n.input) == 3:
                        # Resize version 11 and up (no size input)
                        scales = model.get_initializer(n.input[2])
                    elif len(n.input) == 4:
                        # Resize version 11 and up
                        scales_exists = (model.get_initializer(n.input[2]) is not None) and (
                            len(model.get_initializer(n.input[2])) != 0
                        )
                        sizes_exists = (model.get_initializer(n.input[3]) is not None) and (
                            len(model.get_initializer(n.input[3])) != 0
                        )
                        assert scales_exists ^ sizes_exists, (
                            "%s: Either scales or the target output size must "
                            "be specified. Specifying both is prohibited." % n.name
                        )
                        if scales_exists:
                            # Scales input
                            scales = model.get_initializer(n.input[2])
                        else:
                            # Convert sizes to scales
                            sizes = model.get_initializer(n.input[3])
                            data_input_size = model.get_tensor_shape(n.input[0])
                            scales = sizes / data_input_size
                in_shape = model.get_tensor_shape(n.input[0])

                dt = model.get_tensor_datatype(n.input[0])
                if not dt.is_integer():
                    log.warning(f"{n.name}: Input not int. Can't infer UpsampleNearestNeighbour.")
                    continue

                if model.get_tensor_layout(n.input[0]) != DataLayout.NHWC:
                    log.warning(f"{n.name}: Input not NHWC. Can't infer UpsampleNearestNeighbour.")
                    continue

                # Check that the parameters are okay
                assert mode == "nearest", (
                    "%s: Upsampling is only supported for the mode nearest." % n.name
                )
                assert len(in_shape) == 4, "Upsampling is only supported for 4D inputs."
                assert scales.shape == (4,), (
                    "%s: Upsampling is only supported for 4D scales." % n.name
                )
                assert (scales >= 1).all(), (
                    n.name + ": Upsampling is only supported for scales "
                    "which are larger or equal 1 in all dimensions."
                )

                # Assumes nhwc layout for scales and input
                assert scales[0] == scales[3] == 1, (
                    n.name + ": Upsampling is only supported for scales with "
                    "the first and last dimensions being 1 in NHWC."
                )

                # Extract information for HW node
                HI = in_shape[1]
                WI = in_shape[2]
                HO = int(round(HI * scales[1]))
                WO = int(round(WI * scales[2]))
                NumChannels = in_shape[-1]
                batchSize = in_shape[0]
                inputDataType = dt.name

                # Insert the HWCustomOp node
                Upsample_HW_node = helper.make_node(
                    "UpsampleNearestNeighbour",
                    [n.input[0]],
                    [n.output[0]],
                    domain="finn.custom_op.fpgadataflow",
                    backend="fpgadataflow",
                    SIMD=1,
                    HO=HO,
                    WO=WO,
                    HI=HI,
                    WI=WI,
                    NumChannels=NumChannels,
                    inputDataType=inputDataType,
                    batchSize=batchSize,
                    name="UpsampleNearestNeighbour_" + n.name,
                    cpp_interface="hls_vector",
                    hls_style="freerunning",
                )

                # Remove the old node
                graph.node.insert(node_ind, Upsample_HW_node)
                # remove old nodes
                graph.node.remove(n)
                graph_modified = True
        return (model, graph_modified)


class InferAddStreamsLayer(Transformation):
    """Convert any Add into a AddStreams HW layer."""

    def apply(self, model):
        """Apply the transformation to convert Add operations into AddStreams HW layers."""
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for node in graph.node:
            node_ind += 1
            if node.op_type == "Add":
                in0 = node.input[0]
                in1 = node.input[1]
                result = node.output[0]
                in0_shape = model.get_tensor_shape(in0)
                in1_shape = model.get_tensor_shape(in1)
                in0_static = not (model.get_initializer(in0) is None)
                in1_static = not (model.get_initializer(in1) is None)

                # skip if different shapes on inputs
                if in0_shape != in1_shape:
                    continue
                # skip if any of inputs have initializers
                # (this node is meant for adding two dynamic streams)
                if in0_static or in1_static:
                    continue

                idt0 = model.get_tensor_datatype(in0)
                idt1 = model.get_tensor_datatype(in1)

                # skip conversion for layers with float input
                if not idt0.is_integer() or not idt1.is_integer():
                    continue

                # check layout and convert if necessary
                in0_layout = model.get_tensor_layout(in0)
                in1_layout = model.get_tensor_layout(in1)
                result_layout = model.get_tensor_layout(result)

                if in0_layout == DataLayout.NCHW:
                    in0 = nchw_to_nhwc(in0, model, node_ind)
                    node_ind += 1
                    in0_shape = model.get_tensor_shape(in0)

                if in1_layout == DataLayout.NCHW:
                    in1 = nchw_to_nhwc(in1, model, node_ind)
                    node_ind += 1
                    in1_shape = model.get_tensor_shape(in1)

                # keep track of where we need to insert the HW Op
                # it has to be ahead of the output transform
                insert_point = node_ind

                if result_layout == DataLayout.NCHW:
                    result = nchw_to_nhwc(result, model, node_ind, reverse=True)
                    node_ind += 1

                # now safe to assume num_channels is size of last dimension
                num_channels = int(in0_shape[-1])
                # create node with no parallelization first
                pe = 1

                # create and insert new AddStreams node
                new_node = helper.make_node(
                    "AddStreams",
                    [in0, in1],
                    [result],
                    domain="finn.custom_op.fpgadataflow",
                    backend="fpgadataflow",
                    NumChannels=num_channels,
                    PE=pe,
                    inputDataTypes=[idt0.name, idt1.name],
                    numInputVectors=in0_shape[:-1],
                    name="AddStreams_" + node.name,
                )
                graph.node.insert(insert_point, new_node)
                # remove old node
                graph.node.remove(node)
                graph_modified = True

        if graph_modified:
            model = model.transform(InferShapes())
            model = model.transform(InferDataTypes())
        return (model, graph_modified)


class InferDuplicateStreamsLayer(Transformation):
    """Insert a DuplicateStreams HW layer for any tensor with fanout >= 2."""

    def apply(self, model):
        """Apply the transformation to insert DuplicateStreams HW layers where needed."""
        graph = model.graph
        node_ind = 0
        graph_modified = False
        # check first if global input is split
        successors = model.find_consumers(graph.input[0].name)
        dt = model.get_tensor_datatype(graph.input[0].name)
        if successors is not None and len(successors) >= 2 and dt.is_integer():
            output_tensor = graph.input[0].name
            n_outputs = len(successors)
            dt = model.get_tensor_datatype(output_tensor)

            # create clone tensors
            out_shape = model.get_tensor_shape(output_tensor)
            out_tensor_clones = []
            for i in range(n_outputs):
                clone = helper.make_tensor_value_info(
                    model.make_new_valueinfo_name(), TensorProto.FLOAT, out_shape
                )
                model.graph.value_info.append(clone)
                out_tensor_clones += [clone.name]

            num_ch = int(out_shape[-1])
            vecs = out_shape[:-1]

            # create node with no parallelization first
            pe = 1

            dup_node = helper.make_node(
                "DuplicateStreams",
                [output_tensor],
                out_tensor_clones,
                domain="finn.custom_op.fpgadataflow",
                backend="fpgadataflow",
                NumChannels=num_ch,
                PE=pe,
                inputDataType=dt.name,
                numInputVectors=vecs,
                NumOutputStreams=n_outputs,
                outFIFODepths=[2] * n_outputs,
                name="DuplicateStreams_" + output_tensor,
            )

            graph.node.insert(0, dup_node)

            # connect successors to out tensor clone
            clone_idx = 0
            for successor in successors:
                for i, succ_input in enumerate(successor.input):
                    if succ_input == output_tensor:
                        successor.input[i] = out_tensor_clones[clone_idx]
                        clone_idx += 1
                        # if one node has multiple connections to the same output
                        # find_direct_successors will return one node per input
                        # so break the inner loop will result in correct behaviour
                        break
            graph_modified = True

        for node in graph.node:
            node_ind += 1
            for output_tensor in node.output:
                successors = model.find_consumers(output_tensor)
                if successors is not None and len(successors) >= 2:
                    n_outputs = len(successors)

                    dt = model.get_tensor_datatype(output_tensor)

                    # skip conversion for layers with float input
                    if not dt.is_integer():
                        continue

                    # create clone tensors
                    out_shape = model.get_tensor_shape(output_tensor)
                    out_tensor_clones = []
                    for i in range(n_outputs):
                        clone = helper.make_tensor_value_info(
                            model.make_new_valueinfo_name(), TensorProto.FLOAT, out_shape
                        )
                        model.graph.value_info.append(clone)
                        out_tensor_clones += [clone.name]

                    num_ch = int(out_shape[-1])
                    vecs = out_shape[:-1]

                    # create node with no parallelization first
                    pe = 1

                    dup_node = helper.make_node(
                        "DuplicateStreams",
                        [output_tensor],
                        out_tensor_clones,
                        domain="finn.custom_op.fpgadataflow",
                        backend="fpgadataflow",
                        NumChannels=num_ch,
                        PE=pe,
                        inputDataType=dt.name,
                        numInputVectors=vecs,
                        NumOutputStreams=n_outputs,
                        outFIFODepths=[2] * n_outputs,
                        name="DuplicateStreams_" + node.name,
                    )

                    graph.node.insert(node_ind, dup_node)

                    # connect successors to out tensor clone
                    clone_idx = 0
                    for successor in successors:
                        for i, succ_input in enumerate(successor.input):
                            if succ_input == output_tensor:
                                successor.input[i] = out_tensor_clones[clone_idx]
                                clone_idx += 1
                                # if one node has multiple connections to the same output
                                # find_direct_successors will return one node per input
                                # so break the inner loop will result in correct behaviour
                                break

                    graph_modified = True

        if graph_modified:
            model = model.transform(SortGraph())
            model = model.transform(InferShapes())
            model = model.transform(InferDataTypes())
        return (model, graph_modified)


class InferChannelwiseLinearLayer(Transformation):
    """Convert any channel-wise Add/Mul into a HW layer."""

    def get_smallest_possible(self, vals):
        """Return smallest (fewest bits) possible DataType that can represent value.

        Prefers unsigned integers where possible.
        """
        vals = np.array(vals, dtype=np.float64)
        for v in vals:
            assert int(v) == v, "Error float value"

        for k in DataType.get_accumulator_dt_cands():
            dt = DataType[k]

            if dt in [DataType["BIPOLAR"], DataType["TERNARY"], DataType["FLOAT32"]]:
                # not currently supported
                continue

            if (dt.min() <= vals).all() and (vals <= dt.max()).all():
                return dt

        log.warning(
            """InferChannelwiseLinearLayer: Output values may not be
        representable with supported data types.
        Setting maximum width data type available.
        This will lead to errors if there are no constrains on the input
        """
        )

        if (0 <= vals).all():
            return DataType["UINT64"]
        else:
            return DataType["INT64"]

    def apply(self, model):
        """Apply transformation to convert channel-wise Add/Mul operations to ChannelwiseOp nodes.

        This transformation identifies Add and Mul operations that operate channel-wise
        and converts them to FINN's custom ChannelwiseOp nodes for hardware acceleration.
        """
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for node in graph.node:
            node_ind += 1
            if node.op_type == "Add" or node.op_type == "Mul":
                # assuming input[0] is dynamic
                ll_input = node.input[0]
                ll_output = node.output[0]
                ll_in_shape = model.get_tensor_shape(ll_input)

                # check if input 1 has an initializer
                ll_const = node.input[1]
                if ll_const is not None:
                    ll_cinit = model.get_initializer(ll_const)
                    if ll_cinit is None:
                        # input 1 is also dynamic
                        continue
                else:
                    continue

                # get number of channels and channel index from input
                ll_in_layout = model.get_tensor_layout(ll_input)
                if ll_in_layout == DataLayout.NHWC or ll_in_layout == DataLayout.NC:
                    ch_index = -1
                    ch = ll_in_shape[-1]
                elif ll_in_layout == DataLayout.NCHW:
                    ch_index = 1
                    ch = ll_in_shape[1]
                else:
                    continue

                # check if the shape of initializer is compatible
                ll_cinit_shape = list(ll_cinit.shape)
                if np.prod(ll_cinit_shape) == 1:
                    log.warning(f"Broadcasting {node.op_type} ({node.name})")
                    ll_cinit = np.full((ch), ll_cinit.flatten()[0])
                elif np.prod(ll_cinit_shape) != ch or ll_cinit_shape[ch_index] != ch:
                    # parameter shape not compatible with Channelwise
                    continue

                # check initializer contains integers as floats
                if not (ll_cinit.astype(np.int32) == ll_cinit).all():
                    continue
                # all initializer conditions are met

                # check inputs
                idt = model.get_tensor_datatype(ll_input)
                if not idt.is_integer():
                    # skip conversion for layers with float input
                    continue

                # check layout of inputs/outputs, and convert if needed
                # check layout and convert if necessary
                if ll_in_layout == DataLayout.NCHW:
                    ll_input = nchw_to_nhwc(ll_input, model, node_ind)
                    node_ind += 1
                    ll_in_shape = model.get_tensor_shape(ll_input)

                # keep track of where we need to insert the HW Op
                # it has to be ahead of the output transform
                insert_point = node_ind
                ll_output_layout = model.get_tensor_layout(ll_output)
                if ll_output_layout == DataLayout.NCHW:
                    ll_output = nchw_to_nhwc(ll_output, model, node_ind, reverse=True)
                    node_ind += 1

                # get parameter data type
                param_min = min(ll_cinit.flatten())
                param_max = max(ll_cinit.flatten())
                pdt = self.get_smallest_possible([param_min, param_max])

                # set function and determine output data type
                if node.op_type == "Add":
                    func = "add"
                    out_min = idt.min() + param_min
                    out_max = idt.max() + param_max
                    odt = self.get_smallest_possible([out_min, out_max])
                elif node.op_type == "Mul":
                    func = "mul"
                    possible_limits = []
                    possible_limits += [idt.min() * param_min]
                    possible_limits += [idt.min() * param_max]
                    possible_limits += [idt.max() * param_min]
                    possible_limits += [idt.max() * param_max]
                    odt = self.get_smallest_possible(possible_limits)

                model.set_initializer(ll_const, ll_cinit.reshape(ch))
                model.set_tensor_datatype(ll_output, odt)

                # create node with no parallelization first
                pe = 1
                assert ch % pe == 0, "Requirement IFC divisable by PE is violated."
                # create and insert node
                new_node = helper.make_node(
                    "ChannelwiseOp",
                    [ll_input, ll_const],
                    [ll_output],
                    domain="finn.custom_op.fpgadataflow",
                    backend="fpgadataflow",
                    Func=func,
                    NumChannels=ch,
                    PE=pe,
                    inputDataType=idt.name,
                    paramDataType=pdt.name,
                    outputDataType=odt.name,
                    numInputVectors=list(ll_in_shape[:-1]),
                    name="ChannelwiseOp_" + node.name,
                )
                graph.node.insert(insert_point, new_node)
                # remove old node
                graph.node.remove(node)
                graph_modified = True

        if graph_modified:
            model = model.transform(InferShapes())
            model = model.transform(InferDataTypes())
        return (model, graph_modified)


class InferLabelSelectLayer(Transformation):
    """Convert any TopK into a LabelSelect HW layer."""

    def apply(self, model):
        """Apply transformation to convert TopK nodes to LabelSelect hardware layers.

        This transformation identifies TopK operations and converts them to FINN's
        custom LabelSelect nodes for hardware acceleration.
        """
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for node in graph.node:
            node_ind += 1
            if node.op_type == "TopK":
                fc_input = node.input[0]
                k_input = node.input[1]
                val_output = node.output[0]
                idx_output = node.output[1]
                fc_in_shape = model.get_tensor_shape(fc_input)

                idt = model.get_tensor_datatype(fc_input)

                # skip conversion for layers with float input
                if not idt.is_integer():
                    continue

                # skip conversion for if value output is connected (not supported)
                if model.find_consumer(val_output) is not None:
                    continue

                num_labels = int(fc_in_shape[-1])
                num_inp_vecs = list(fc_in_shape[:-1])
                # create node with no parallelization first
                pe = 1

                k = model.get_initializer(k_input)[0]

                # create and insert new LabelSelect node
                new_node = helper.make_node(
                    "LabelSelect",
                    [fc_input],
                    [idx_output],
                    domain="finn.custom_op.fpgadataflow",
                    backend="fpgadataflow",
                    Labels=num_labels,
                    PE=pe,
                    K=k,
                    inputDataType=idt.name,
                    numInputVectors=num_inp_vecs,
                    name="LabelSelect_" + node.name,
                )
                graph.node.insert(node_ind, new_node)
                # remove old node
                graph.node.remove(node)
                graph_modified = True

        if graph_modified:
            model = model.transform(InferShapes())
            model = model.transform(InferDataTypes())
        return (model, graph_modified)


class InferGlobalAccPoolLayer(Transformation):
    """Convert any GlobalAveragePool into a GlobalAccPool HW layer and a scalar Mul."""

    def apply(self, model):
        """Apply transformation to infer GlobalAccPool hardware layers."""
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for node in graph.node:
            node_ind += 1
            if node.op_type == "GlobalAveragePool":
                in0 = node.input[0]
                result = node.output[0]
                in0_shape = model.get_tensor_shape(in0)

                idt = model.get_tensor_datatype(in0)

                # skip conversion for layers with float input
                if not idt.is_integer():
                    continue

                # check layout and convert if necessary
                in0_layout = model.get_tensor_layout(in0)
                result_layout = model.get_tensor_layout(result)

                if in0_layout == DataLayout.NCHW:
                    in0 = nchw_to_nhwc(in0, model, node_ind)
                    node_ind += 1
                    in0_shape = model.get_tensor_shape(in0)

                # keep track of where we need to insert the HW Op
                # it has to be ahead of the output transform
                insert_point = node_ind

                if result_layout == DataLayout.NCHW:
                    result = nchw_to_nhwc(result, model, node_ind, reverse=True)
                    node_ind += 1

                num_ch = int(in0_shape[-1])
                vecs = in0_shape[:-1]
                # create node with no parallelization first
                pe = 1

                # create an additional tensor of the same shape and layout as result
                out_shape = model.get_tensor_shape(result)
                pool_out = helper.make_tensor_value_info(
                    model.make_new_valueinfo_name(), TensorProto.FLOAT, out_shape
                )
                model.graph.value_info.append(pool_out)
                pool_out = pool_out.name
                model.set_tensor_layout(pool_out, model.get_tensor_layout(result))

                new_pool = helper.make_node(
                    "GlobalAccPool",
                    [in0],
                    [pool_out],
                    domain="finn.custom_op.fpgadataflow",
                    backend="fpgadataflow",
                    NumChannels=num_ch,
                    PE=pe,
                    inputDataType=idt.name,
                    numInputVectors=vecs,
                    name="GlobalAccPool_" + node.name,
                )

                mul_value = helper.make_tensor_value_info(
                    model.make_new_valueinfo_name(), TensorProto.FLOAT, [1]
                )
                model.graph.value_info.append(mul_value)
                model.set_initializer(
                    mul_value.name, np.array(1 / (vecs[1] * vecs[2]), dtype=np.float32)
                )
                new_mul = helper.make_node(
                    "Mul",
                    [pool_out, mul_value.name],
                    [result],
                )
                graph.node.insert(insert_point, new_pool)
                graph.node.insert(insert_point + 1, new_mul)
                node_ind += 1
                # remove old node
                graph.node.remove(node)
                graph_modified = True

        if graph_modified:
            model = model.transform(InferShapes())
            model = model.transform(InferDataTypes())
        return (model, graph_modified)


class InferPool(Transformation):
    """If kernel_shape > strides, replace Pool layer with Im2col + pool combination.

    When kernel_shape > strides, replaces Pool layer with Im2col followed by
    pool (with kernel_shape == strides), plus Transpose layers to keep the
    original data layout.
    """

    def apply(self, model):
        """Apply transformation to convert Pool operations with kernel_shape > strides."""
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for node in graph.node:
            node_ind += 1
            if node.op_type in ["MaxPool", "QuantAvgPool2d", "MaxPoolNHWC"]:
                node_input = node.input[0]
                ishape = model.get_tensor_shape(node_input)
                node_output = node.output[0]
                idt = model.get_tensor_datatype(node_input)
                oshape = model.get_tensor_shape(node_output)
                # only support 4D input tensors (1D convs need extra dummy dim)
                if len(ishape) != 4:
                    continue

                # extract pool parameters
                if node.op_type == "MaxPool":
                    kh, kw = list(get_by_name(node.attribute, "kernel_shape").ints)
                    sh, sw = list(get_by_name(node.attribute, "strides").ints)
                    dlayout = "NCHW"
                elif node.op_type == "QuantAvgPool2d":
                    inst = getCustomOp(node)
                    # QuantAvgPool2d has a single scalar attribute
                    # for kernel size and stride (implicit square)
                    kh = kw = inst.get_nodeattr("kernel")
                    sh = sw = inst.get_nodeattr("stride")
                    dlayout = inst.get_nodeattr("data_layout")
                elif node.op_type == "MaxPoolNHWC":
                    inst = getCustomOp(node)
                    kh, kw = inst.get_nodeattr("kernel_shape")
                    sh, sw = inst.get_nodeattr("strides")
                    dlayout = "NHWC"
                try:
                    pad = list(get_by_name(node.attribute, "pads").ints)
                except AttributeError:
                    pad = [0, 0, 0, 0]

                if not idt.is_integer():
                    continue

                if (kh < sh) or (kw < sw):
                    # TODO check/implement swg support
                    continue

                odt = model.get_tensor_datatype(node_output)

                if dlayout == "NCHW":
                    _, ifm_ch, ifm_h, ifm_w = ishape
                    _, ofm_ch, ofm_h, ofm_w = oshape
                elif dlayout == "NHWC":
                    _, ifm_h, ifm_w, ifm_ch = ishape
                    _, ofm_h, ofm_w, ofm_ch = oshape
                else:
                    raise Exception("Unknown dlayout: " + str(dlayout))

                # if data layout NCHW, we need transpose nodes surrounding
                # the hw layer
                if dlayout == "NCHW":
                    # create new intermediate values
                    inp_trans_out = helper.make_tensor_value_info(
                        model.make_new_valueinfo_name(),
                        TensorProto.FLOAT,
                        (1, ifm_h, ifm_w, ifm_ch),  # NHWC
                    )
                    graph.value_info.append(inp_trans_out)
                    inp_trans_out = inp_trans_out.name
                    model.set_tensor_datatype(inp_trans_out, idt)

                    pool_output = helper.make_tensor_value_info(
                        model.make_new_valueinfo_name(),
                        TensorProto.FLOAT,
                        (1, ofm_h, ofm_w, ofm_ch),
                    )
                    graph.value_info.append(pool_output)
                    pool_output = pool_output.name

                im2col_out = helper.make_tensor_value_info(
                    model.make_new_valueinfo_name(),
                    TensorProto.FLOAT,
                    (1, ofm_h, ofm_w, ifm_ch * kh * kw),
                )
                graph.value_info.append(im2col_out)
                im2col_out = im2col_out.name
                model.set_tensor_datatype(im2col_out, idt)

                # create new nodes
                if dlayout == "NCHW":
                    # NCHW -> NHWC
                    inp_trans_node = helper.make_node(
                        "Transpose", [node_input], [inp_trans_out], perm=[0, 2, 3, 1]
                    )
                    im2col_in = inp_trans_out
                else:
                    im2col_in = node_input
                    pool_output = node_output

                accum_bits = 0
                pool_size_param = 0  # will be overridden if neededs
                pad_value = 0
                if node.op_type in ["MaxPool", "MaxPoolNHWC"]:
                    pool_fxn = "MaxPool"
                    odt = idt
                    pad_value = idt.min()
                elif node.op_type == "QuantAvgPool2d":
                    assert odt.is_integer(), """Output data type for QuantAvgPool2d
                    needs to be integer"""
                    assert all(x == 0 for x in pad), "Padding is not supported for QuantAvgPool2d"
                    inst = getCustomOp(node)
                    pool_fxn = "QuantAvgPool"
                    pool_size_param = inst.get_shifts()
                    accum_bits = inst.get_accum_size()

                else:
                    raise Exception(
                        "pad_value and pool_fxn not configured for {}".format(node.op_type)
                    )

                # format input tensor
                im2col_node = helper.make_node(
                    "Im2Col",
                    [im2col_in],
                    [im2col_out],
                    domain="qonnx.custom_op.general",
                    stride=[sh, sw],
                    kernel_size=[kh, kw],
                    pad_amount=pad,
                    pad_value=pad_value,
                    depthwise=1,
                    input_shape="(1,{},{},{})".format(ifm_h, ifm_w, ifm_ch),
                    name="Im2Col_" + node.name,
                )

                # Warning PE has to be equal to ifm_ch until Im2Col is replaced by
                # ConvolutionInputGenerator with depthwise=1.
                # For other settings the output will be incorrect due to incorrect input
                # data layout
                pool_node = helper.make_node(
                    "Pool",
                    [im2col_out],
                    [pool_output],
                    domain="finn.custom_op.fpgadataflow",
                    backend="fpgadataflow",
                    InputDataType=idt.name,
                    OutputDataType=odt.name,
                    Channels=ifm_ch,
                    PE=ifm_ch,
                    KernelSize=[kh, kw],
                    Function=pool_fxn,
                    OutImgDims=[ofm_h, ofm_w],
                    AccumBits=accum_bits,
                    Size=pool_size_param,
                    BatchSize=1,
                    cpp_interface="hls_vector",
                    name="Pool_" + node.name,
                )

                if dlayout == "NCHW":
                    # NHWC -> NCHW
                    out_trans_node = helper.make_node(
                        "Transpose", [pool_output], [node_output], perm=[0, 3, 1, 2]
                    )

                # insert nodes where the conv is to preserve topological ordering
                if dlayout == "NCHW":
                    graph.node.insert(node_ind, inp_trans_node)
                    graph.node.insert(node_ind + 1, im2col_node)
                    graph.node.insert(node_ind + 2, pool_node)
                    graph.node.insert(node_ind + 3, out_trans_node)
                else:
                    graph.node.insert(node_ind, im2col_node)
                    graph.node.insert(node_ind + 1, pool_node)
                # remove old node
                graph.node.remove(node)
                graph_modified = True

        if graph_modified:
            model = model.transform(InferShapes())
            model = model.transform(InferDataTypes())
        return (model, graph_modified)


class InferLookupLayer(Transformation):
    """Convert Gather nodes with constant op0 into Lookup HW layers."""

    def apply(self, model):
        """Apply transformation to convert Gather operations to Lookup hardware layers.

        This transformation identifies Gather operations with constant first operand
        and converts them to FINN's custom Lookup nodes for hardware acceleration.
        """
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for node in graph.node:
            node_ind += 1
            if node.op_type == "Gather":
                emb_name = node.input[0]
                embs = model.get_initializer(emb_name)
                axis = get_by_name(node.attribute, "axis")
                # skip conversion if input0 is not constant
                if embs is None:
                    continue
                # skip conversion if axis != 0
                if axis is not None and axis.i != 0:
                    continue
                ind_name = node.input[1]
                ind_dtype = model.get_tensor_datatype(ind_name)
                emb_dtype = model.get_tensor_datatype(emb_name)
                # skip conversion if inputs are not unsigned integers
                if (not ind_dtype.is_integer()) or ind_dtype.signed():
                    continue
                num_embs, emb_dim = embs.shape
                out_name = node.output[0]
                ishape = model.get_tensor_shape(node.input[1])
                # create and insert new Lookup node
                new_node = helper.make_node(
                    "Lookup",
                    [ind_name, emb_name],
                    [out_name],
                    domain="finn.custom_op.fpgadataflow",
                    backend="fpgadataflow",
                    name="Lookup_" + node.name,
                    NumEmbeddings=num_embs,
                    EmbeddingDim=emb_dim,
                    EmbeddingType=emb_dtype.name,
                    InputType=ind_dtype.name,
                    InputShape=list(ishape),
                )
                graph.node.insert(node_ind, new_node)
                # remove old node
                graph.node.remove(node)
                graph_modified = True

        if graph_modified:
            model = model.transform(InferShapes())
            model = model.transform(InferDataTypes())
        return (model, graph_modified)


class InferConcatLayer(Transformation):
    """Convert suitable Concat nodes (operating on last/-1 axis) into StreamingConcat HW layers."""

    def apply(self, model):
        """Apply transformation to convert Concat operations to StreamingConcat hardware layers.

        This transformation identifies Concat operations operating on the last axis
        and converts them to FINN's custom StreamingConcat nodes.
        """
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for node in graph.node:
            node_ind += 1
            if node.op_type == "Concat":
                ishape = model.get_tensor_shape(node.input[0])
                axis = get_by_name(node.attribute, "axis")
                if (axis is None) or (ishape is None):
                    continue
                axis = axis.i
                last_axis = len(ishape) - 1
                # skip conversion if not using last axis
                if (axis != -1) and (axis != last_axis):
                    continue
                # check datatype coherence
                if any([model.get_tensor_datatype(x) is None for x in node.input]):
                    log.warning(
                        "Inputs with undefined datatype detected, skipping InferConcatLayer()"
                    )
                    continue
                # skip conversion if any inputs are static
                any_static = any([model.get_initializer(x) is not None for x in node.input])
                if any_static:
                    continue
                # skip conversion if inputs are not integers
                all_integer = all([model.get_tensor_datatype(x).is_integer() for x in node.input])
                if not all_integer:
                    log.warning(
                        "Inputs with non-integer datatype detected, skipping InferConcatLayer()"
                    )
                    continue
                # ready for conversion
                channels_per_stream = [model.get_tensor_shape(x)[-1] for x in node.input]
                inp_vec = list(model.get_tensor_shape(node.input[0])[:-1])
                new_node = helper.make_node(
                    "StreamingConcat",
                    node.input,
                    node.output,
                    domain="finn.custom_op.fpgadataflow",
                    backend="fpgadataflow",
                    name="StreamingConcat_" + node.name,
                    SIMD=1,
                    ChannelsPerStream=channels_per_stream,
                    inputDataTypes=[model.get_tensor_datatype(x).name for x in node.input],
                    numInputVectors=inp_vec,
                    inFIFODepths=[2] * len(node.input),
                    cpp_interface="hls_vector",
                    hls_style="freerunning",
                )
                graph.node.insert(node_ind, new_node)
                # remove old node
                graph.node.remove(node)
                graph_modified = True

        if graph_modified:
            model = model.transform(InferShapes())
            model = model.transform(InferDataTypes())
        return (model, graph_modified)


class InferSplitLayer(Transformation):
    """Convert suitable Split nodes (operating on last/-1 axis) into StreamingSplit HW layers."""

    def apply(self, model):
        """Apply transformation to convert Split operations to StreamingSplit hardware layers.

        This transformation identifies Split operations operating on the last axis
        and converts them to FINN's custom StreamingSplit nodes.
        """
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for node in graph.node:
            node_ind += 1
            if node.op_type == "Split":
                split_param = node.input[1]
                if model.get_initializer(split_param) is None:
                    log.warning("Split param not constant, skipping InferSplitLayer()")
                    continue
                ishape = model.get_tensor_shape(node.input[0])
                axis = get_by_name(node.attribute, "axis")
                if (axis is None) or (ishape is None):
                    continue
                axis = axis.i
                last_axis = len(ishape) - 1
                # skip conversion if not using last axis
                if (axis != -1) and (axis != last_axis):
                    log.warning(
                        "StreamingSplit supports only last axis, skipping InferSplitLayer()"
                    )
                    continue
                # only one input allowed (two including split_param)
                if len(node.input) != 2:
                    log.warning("Only one input allowed, skipping InferSplitLayer()")
                    continue
                # skip conversion if the input is static
                if model.get_initializer(node.input[0]) is not None:
                    log.warning("Static input detected, skipping InferSplitLayer()")
                    continue
                # skip conversion if inputs are not integers
                if not model.get_tensor_datatype(node.input[0]).is_integer():
                    log.warning("Non-integer input detected, skipping InferSplitLayer()")
                    continue
                # ready for conversion
                channels_per_stream = [model.get_tensor_shape(x)[-1] for x in node.output]
                inp_vec = list(model.get_tensor_shape(node.input[0])[:-1])
                # when creating the fpgadataflow node we remove the second parameter input
                new_node = helper.make_node(
                    "StreamingSplit",
                    [node.input[0]],
                    node.output,
                    domain="finn.custom_op.fpgadataflow",
                    backend="fpgadataflow",
                    name="StreamingSplit_" + node.name,
                    SIMD=1,
                    cpp_interface="hls_vector",
                    hls_style="freerunning",
                    ChannelsPerStream=channels_per_stream,
                    inputDataType=model.get_tensor_datatype(node.input[0]).name,
                    numInputVectors=inp_vec,
                    outFIFODepths=[2] * len(node.output),
                )
                graph.node.insert(node_ind, new_node)
                # remove old node
                graph.node.remove(node)
                graph_modified = True

        if graph_modified:
            model = model.transform(InferShapes())
            model = model.transform(InferDataTypes())
        return (model, graph_modified)


class InferStreamingEltwise(Transformation):
    """Convert eltwise Add, Sub or Sub -> Abs to StreamingEltwise layer.

    Converts element-wise operations with AddEltwise, SubEltwise or AbsDiffEltwise op.
    """

    def apply(self, model):
        """Apply transformation to convert element-wise operations to StreamingEltwise layers."""
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for node in graph.node:
            node_ind += 1
            if node.op_type in ["Sub", "Add"]:
                in0 = node.input[0]
                in1 = node.input[1]
                result = node.output[0]
                in0_shape = model.get_tensor_shape(in0)
                in1_shape = model.get_tensor_shape(in1)
                in0_static = not (model.get_initializer(in0) is None)
                in1_static = not (model.get_initializer(in1) is None)

                # skip if different shapes on inputs
                if in0_shape != in1_shape:
                    continue
                # skip if any of inputs have initializers
                # (this node is meant for two dynamic streams)
                if in0_static or in1_static:
                    continue

                idt0 = model.get_tensor_datatype(in0)
                idt1 = model.get_tensor_datatype(in1)

                # skip conversion for layers with float input
                if not (idt0.is_integer() and idt1.is_integer()):
                    continue

                eltwiseOp = node.op_type
                nodes_to_remove = [node]
                if node.op_type == "Sub":
                    # look for a downstream Abs node
                    res_consumer = model.find_consumer(result)
                    if (res_consumer is not None) and (res_consumer.op_type == "Abs"):
                        eltwiseOp = "AbsDiff"
                        result = res_consumer.output[0]
                        nodes_to_remove.append(res_consumer)

                # check layout and convert if necessary
                in0_layout = model.get_tensor_layout(in0)
                in1_layout = model.get_tensor_layout(in1)
                result_layout = model.get_tensor_layout(result)

                if in0_layout == DataLayout.NCHW:
                    in0 = nchw_to_nhwc(in0, model, node_ind)
                    node_ind += 1
                    in0_shape = model.get_tensor_shape(in0)

                if in1_layout == DataLayout.NCHW:
                    in1 = nchw_to_nhwc(in1, model, node_ind)
                    node_ind += 1
                    in1_shape = model.get_tensor_shape(in1)

                # keep track of where we need to insert the HW Op
                # it has to be ahead of the output transform
                insert_point = node_ind

                if result_layout == DataLayout.NCHW:
                    result = nchw_to_nhwc(result, model, node_ind, reverse=True)
                    node_ind += 1

                # now safe to assume num_channels is size of last dimension
                num_channels = int(in0_shape[-1])
                # create node with no parallelization first
                pe = 1

                # create and insert new Eltwise node
                new_node = helper.make_node(
                    "StreamingEltwise",
                    [in0, in1],
                    [result],
                    domain="finn.custom_op.fpgadataflow",
                    backend="fpgadataflow",
                    NumChannels=num_channels,
                    PE=pe,
                    inputDataType0=idt0.name,
                    inputDataType1=idt1.name,
                    eltwiseOp=eltwiseOp,
                    numInputVectors=in0_shape[:-1],
                    name="StreamingEltwise_" + node.name,
                )
                graph.node.insert(insert_point, new_node)
                # remove old nodes
                for nd in nodes_to_remove:
                    graph.node.remove(nd)
                graph_modified = True

        return (model, graph_modified)


class InferBinaryMatrixVectorActivation(Transformation):
    """Convert XnorPopcountMatMul layers to MatrixVectorActivation layers.

    Any immediately following MultiThreshold layers will also be absorbed into the MVTU.
    """

    def __init__(self):
        """Initialize the transformation."""
        super().__init__()

    def apply(self, model):
        """Apply transformation to convert XnorPopcountMatMul to MVAU nodes.

        This transformation identifies XnorPopcountMatMul operations and converts them
        to FINN's custom MVAU (Matrix Vector Activation Unit) nodes, potentially
        absorbing following MultiThreshold layers.
        """
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for n in graph.node:
            node_ind += 1
            if n.op_type == "XnorPopcountMatMul":
                mm_input = n.input[0]
                mm_weight = n.input[1]
                mm_output = n.output[0]
                mm_in_shape = model.get_tensor_shape(mm_input)
                mm_out_shape = model.get_tensor_shape(mm_output)
                assert model.get_tensor_datatype(mm_input) == DataType["BINARY"], (
                    n.name
                    + """: First
                input for xnorpopcount is not Wset to FINN DataType BINARY."""
                )
                assert model.get_tensor_datatype(mm_weight) == DataType["BINARY"], (
                    n.name
                    + """: Second
                input (weights) for xnorpopcount is not set to FINN DataType BINARY."""
                )
                idt = DataType["BINARY"]
                wdt = DataType["BINARY"]
                mm_output = n.output[0]
                W = model.get_initializer(mm_weight)
                # extract weight shape, note that ONNX and finn-hlslib
                # make different assumptions about dim order here
                # ONNX assumes W has (in, out) shape
                # finn-hlslib assumes W has (out, in) shape
                mh = int(W.shape[1])
                mw = int(W.shape[0])
                # create node with no parallelization first
                pe = 1
                simd = 1
                wmem = mw * mh // (pe * simd)
                assert mw * mh == wmem * pe * simd, (
                    n.name
                    + """: Requirement (MW * MH) divisiable by
                (WMEM * PE * SIMD) is violated."""
                )
                # see if we have any following thresholds
                consumers = model.find_consumers(mm_output)
                # Only a single consumer node can be absorbed. Absorbing one
                # branch of a forking matmul would lead to detached nodes
                # breaking the graph.
                consumer = consumers[0] if len(consumers) == 1 else None
                if consumer is not None and consumer.op_type == "MultiThreshold":
                    # TODO ensure integer thresholds?
                    # create MVTU (i.e. including activation)
                    mt_output = consumer.output[0]
                    mt_out_shape = model.get_tensor_shape(mt_output)
                    mt_thres = consumer.input[1]
                    T = model.get_initializer(mt_thres)
                    assert T.shape[0] == 1 or T.shape[0] == mh, (
                        consumer.name
                        + """: First dimension of
                    thresholds neither 1 nor MH."""
                    )
                    odt = model.get_tensor_datatype(mt_output)
                    if odt.bitwidth() == 1:
                        # covers both bipolar and binary
                        actval = 0
                    else:
                        actval = odt.min()
                    model.set_tensor_shape(mm_input, mm_in_shape)
                    model.set_tensor_shape(mt_output, mt_out_shape)
                    # create and insert new MatrixVectorActivation node
                    new_node = helper.make_node(
                        "MVAU",
                        [mm_input, mm_weight, mt_thres],
                        [mt_output],
                        domain="finn.custom_op.fpgadataflow",
                        backend="fpgadataflow",
                        MW=mw,
                        MH=mh,
                        SIMD=simd,
                        PE=pe,
                        inputDataType=idt.name,
                        weightDataType=wdt.name,
                        outputDataType=odt.name,
                        ActVal=actval,
                        binaryXnorMode=1,
                        noActivation=0,
                        numInputVectors=list(mm_in_shape[:-1]),
                        name=n.name,
                    )
                    graph.node.insert(node_ind, new_node)
                    # remove old nodes
                    graph.node.remove(n)
                    graph.node.remove(consumer)
                    graph_modified = True
                else:
                    # no activation, matmul only
                    odt = model.get_tensor_datatype(mm_output)
                    model.set_tensor_shape(mm_input, mm_in_shape)
                    model.set_tensor_shape(mm_output, mm_out_shape)
                    # create and insert new MatrixVectorActivation node
                    new_node = helper.make_node(
                        "MVAU",
                        [mm_input, mm_weight],
                        [mm_output],
                        domain="finn.custom_op.fpgadataflow",
                        backend="fpgadataflow",
                        MW=mw,
                        MH=mh,
                        SIMD=simd,
                        PE=pe,
                        inputDataType=idt.name,
                        weightDataType=wdt.name,
                        outputDataType=odt.name,
                        ActVal=0,
                        binaryXnorMode=1,
                        noActivation=1,
                        numInputVectors=list(mm_in_shape[:-1]),
                        name=n.name,
                    )
                    graph.node.insert(node_ind, new_node)
                    # remove old node
                    graph.node.remove(n)
                    graph_modified = True
        if graph_modified:
            model = model.transform(InferShapes())
            model = model.transform(InferDataTypes())
        return (model, graph_modified)


class InferQuantizedMatrixVectorActivation(Transformation):
    """Convert MatMul layers with quantized inputs and weights to
    MatrixVectorActivation layers."""

    def apply(self, model):
        """Apply transformation to convert MatMul to MVAU nodes."""
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for n in graph.node:
            node_ind += 1
            if n.op_type == "MatMul" and model.get_tensor_sparsity(n.input[1]) is None:
                mm_input = n.input[0]
                mm_weight = n.input[1]
                mm_output = n.output[0]
                mm_in_shape = model.get_tensor_shape(mm_input)
                mm_out_shape = model.get_tensor_shape(mm_output)
                idt = model.get_tensor_datatype(mm_input)
                wdt = model.get_tensor_datatype(mm_weight)
                W = model.get_initializer(mm_weight)
                if idt.is_integer() and wdt.is_integer():
                    # extract weight shape, note that ONNX and finn-hlslib
                    # make different assumptions about dim order here
                    # ONNX assumes W has (in, out) shape
                    # finn-hlslib assumes W has (out, in) shape
                    if W is None:
                        # dynamic
                        mm_dyn_shape = model.get_tensor_shape(mm_weight)
                        mh = int(mm_dyn_shape[-1])
                        mw = int(mm_dyn_shape[-2])
                    else:
                        # static
                        mh = int(W.shape[1])
                        mw = int(W.shape[0])
                    # create node with no parallelization first
                    pe = 1
                    simd = 1
                    wmem = mw * mh // (pe * simd)
                    assert mw * mh == wmem * pe * simd, (
                        n.name
                        + """: Requirement (MW * MH) divisible by
                    (WMEM * PE * SIMD) is violated."""
                    )
                    # see if we have any following thresholds
                    consumers = model.find_consumers(mm_output)
                    # Only a single consumer node can be absorbed. Absorbing one
                    # branch of a forking matmul would lead to detached nodes
                    # breaking the graph.
                    consumer = consumers[0] if len(consumers) == 1 else None
                    if consumer is not None and consumer.op_type == "MultiThreshold":
                        # TODO ensure integer thresholds?
                        # create MVTU (i.e. including activation)
                        mt_output = consumer.output[0]
                        mt_out_shape = model.get_tensor_shape(mt_output)
                        mt_thres = consumer.input[1]
                        T = model.get_initializer(mt_thres)
                        assert T.shape[0] == 1 or T.shape[0] == mh, (
                            consumer.name
                            + """: First dimension of
                        thresholds neither 1 nor MH."""
                        )
                        odt = model.get_tensor_datatype(mt_output)
                        scale = getCustomOp(consumer).get_nodeattr("out_scale")
                        actval = getCustomOp(consumer).get_nodeattr("out_bias")
                        assert int(actval) == actval, (
                            consumer.name + ": out_bias must be integer for HLS conversion."
                        )
                        actval = int(actval)
                        odt_is_bipolar = odt == DataType["BIPOLAR"]
                        bipolar_ok = odt_is_bipolar and (scale == 2.0) and (actval == -1)
                        assert scale == 1.0 or bipolar_ok, (
                            consumer.name + ": out_scale=1 or bipolar output needed for conversion."
                        )
                        assert (not odt.signed()) or (actval < 0), (
                            consumer.name + ": Signed output requres actval < 0"
                        )
                        model.set_tensor_shape(mm_input, mm_in_shape)
                        model.set_tensor_shape(mt_output, mt_out_shape)
                        if bipolar_ok:
                            # remove bias for bipolar, since
                            # binary->bipolar is achieved by reinterpretation
                            actval = 0
                        # create and insert new MatrixVectorActivation node
                        new_node = helper.make_node(
                            "MVAU",
                            [mm_input, mm_weight, mt_thres],
                            [mt_output],
                            domain="finn.custom_op.fpgadataflow",
                            backend="fpgadataflow",
                            MW=mw,
                            MH=mh,
                            SIMD=simd,
                            PE=pe,
                            inputDataType=idt.name,
                            weightDataType=wdt.name,
                            outputDataType=odt.name,
                            ActVal=actval,
                            binaryXnorMode=0,
                            noActivation=0,
                            numInputVectors=list(mm_in_shape[:-1]),
                            name="MVAU_" + n.name,
                            dynamic_input=W is None,
                            inFIFODepths=[2, 2] if W is None else [2],
                        )
                        graph.node.insert(node_ind, new_node)
                        # remove old nodes
                        graph.node.remove(n)
                        graph.node.remove(consumer)
                        graph_modified = True
                    else:
                        # no activation, matmul only
                        odt = model.get_tensor_datatype(mm_output)
                        model.set_tensor_shape(mm_input, mm_in_shape)
                        model.set_tensor_shape(mm_output, mm_out_shape)
                        # create and insert new MatrixVectorActivation node
                        new_node = helper.make_node(
                            "MVAU",
                            [mm_input, mm_weight],
                            [mm_output],
                            domain="finn.custom_op.fpgadataflow",
                            backend="fpgadataflow",
                            MW=mw,
                            MH=mh,
                            SIMD=simd,
                            PE=pe,
                            inputDataType=idt.name,
                            weightDataType=wdt.name,
                            outputDataType=odt.name,
                            ActVal=0,
                            binaryXnorMode=0,
                            noActivation=1,
                            numInputVectors=list(mm_in_shape[:-1]),
                            name="MVAU_" + n.name,
                            dynamic_input=W is None,
                            inFIFODepths=[2, 2] if W is None else [2],
                        )
                        graph.node.insert(node_ind, new_node)
                        # remove old node
                        graph.node.remove(n)
                        graph_modified = True
        if graph_modified:
            model = model.transform(InferShapes())
            model = model.transform(InferDataTypes())
        return (model, graph_modified)


class InferVectorVectorActivation(Transformation):
    """Convert MatMul layers to VectorVectorActivation layers for depthwise convolutions.

    Converts MatMul layers with quantized inputs and weights to VectorVectorActivation
    layers, if the sparsity annotation of the weight matrix indicates that the MatMul
    layer belongs to a depthwise convolution. Any immediately following MultiThreshold
    layers will also be absorbed into the VVAU.
    """

    def __init__(self):
        """Initialize the transformation."""
        super().__init__()

    def apply(self, model):
        """Apply transformation to convert MatMul to VVAU nodes for depthwise convolutions."""
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for n in graph.node:
            node_ind += 1
            if n.op_type == "MatMul" and model.get_tensor_sparsity(n.input[1]) is not None:
                sparsity = model.get_tensor_sparsity(n.input[1])
                try:
                    k_h, k_w = sparsity["dw"]["kernel_shape"]
                except KeyError:
                    raise Exception(
                        n.name
                        + """: sparsity annotation doesn't indicate that MatMul
                        belongs to a depthwise convolution."""
                    )

                mm_input = n.input[0]
                mm_weight = n.input[1]
                mm_output = n.output[0]
                mm_in_shape = model.get_tensor_shape(mm_input)
                mm_out_shape = model.get_tensor_shape(mm_output)
                idt = model.get_tensor_datatype(mm_input)
                wdt = model.get_tensor_datatype(mm_weight)
                if idt.is_integer() and wdt.is_integer():
                    mm_output = n.output[0]
                    W = model.get_initializer(mm_weight)
                    # infer dense weight tensor from sparse weight matrix
                    # kernel size (k_h, k_w) which was extracted above and the value of
                    # the channels is used.
                    # the weight matrix has a shape of (k_h * k_w * Channels, Channels)
                    # we need to reverse the creation of the sparse weight matrix
                    # to achieve a weight tensor of shape (Channels, 1, k_h, k_w)
                    channels = int(W.shape[1])
                    # transpose to achieve a shape of (k_h * k_w * Channels, Channels)
                    W = W.T
                    # reshape to (Channels, k_h, k_w, Channels) to transpose afterwards
                    # to (Channels, Channels, k_h, k_w)
                    W = W.reshape(channels, k_h, k_w, channels)
                    W = W.transpose(0, 3, 1, 2)
                    # now we can extract the values using a for loop over the channels
                    # and fill a zero numpy array in the correct shape
                    w_tensor = np.zeros((channels, 1, k_h, k_w), dtype=np.float32)
                    for ch in range(channels):
                        w_tensor[ch][0] = W[ch][ch]
                    model.set_initializer(mm_weight, w_tensor)
                    model.set_tensor_shape(mm_weight, (channels, 1, k_h, k_w))
                    # create node with pe=channels as default
                    pe = channels
                    # see if we have any following thresholds
                    consumers = model.find_consumers(mm_output)
                    # Only a single consumer node can be absorbed. Absorbing one
                    # branch of a forking matmul would lead to detached nodes
                    # breaking the graph.
                    consumer = consumers[0] if len(consumers) == 1 else None
                    if consumer is not None and consumer.op_type == "MultiThreshold":
                        # create VVAU (i.e. including activation)
                        mt_output = consumer.output[0]
                        mt_out_shape = model.get_tensor_shape(mt_output)
                        mt_thres = consumer.input[1]
                        T = model.get_initializer(mt_thres)
                        assert T.shape[0] == 1 or T.shape[0] == channels, (
                            consumer.name
                            + """: First dimension of
                        thresholds neither 1 nor Channels."""
                        )
                        odt = model.get_tensor_datatype(mt_output)
                        scale = getCustomOp(consumer).get_nodeattr("out_scale")
                        assert scale == 1.0, (
                            consumer.name + ": out_scale must be equal to 1.0 for HLS conversion."
                        )
                        actval = getCustomOp(consumer).get_nodeattr("out_bias")
                        assert int(actval) == actval, (
                            consumer.name + ": out_bias must be integer for HLS conversion."
                        )
                        actval = int(actval)
                        assert (not odt.signed()) or (actval < 0), (
                            consumer.name + ": Signed output requres actval < 0"
                        )
                        model.set_tensor_shape(mm_input, mm_in_shape)
                        model.set_tensor_shape(mt_output, mt_out_shape)
                        # create and insert new VectorVectorActivation node
                        new_node = helper.make_node(
                            "VVAU",
                            [mm_input, mm_weight, mt_thres],
                            [mt_output],
                            domain="finn.custom_op.fpgadataflow",
                            backend="fpgadataflow",
                            PE=pe,
                            Dim=[mm_in_shape[1], mm_in_shape[2]],
                            Channels=channels,
                            Kernel=[k_h, k_w],
                            inputDataType=idt.name,
                            weightDataType=wdt.name,
                            outputDataType=odt.name,
                            ActVal=actval,
                            noActivation=0,
                            name="VVAU_" + n.name,
                        )
                        graph.node.insert(node_ind, new_node)
                        # remove old nodes
                        graph.node.remove(n)
                        graph.node.remove(consumer)
                        graph_modified = True
                    else:
                        # no activation, matmul only
                        odt = model.get_tensor_datatype(mm_output)
                        model.set_tensor_shape(mm_input, mm_in_shape)
                        model.set_tensor_shape(mm_output, mm_out_shape)
                        # create and insert new VVAU node
                        new_node = helper.make_node(
                            "VVAU",
                            [mm_input, mm_weight],
                            [mm_output],
                            domain="finn.custom_op.fpgadataflow",
                            backend="fpgadataflow",
                            PE=pe,
                            Dim=[mm_in_shape[1], mm_in_shape[2]],
                            Channels=channels,
                            Kernel=[k_h, k_w],
                            inputDataType=idt.name,
                            weightDataType=wdt.name,
                            outputDataType=odt.name,
                            ActVal=0,
                            noActivation=1,
                            name="VVAU_" + n.name,
                        )
                        graph.node.insert(node_ind, new_node)
                        # remove old node
                        graph.node.remove(n)
                        graph_modified = True
        if graph_modified:
            model = model.transform(InferShapes())
            model = model.transform(InferDataTypes())
        return (model, graph_modified)


def lift_to_rank1(name: str, model: ModelWrapper):
    """Lift scalar to rank-1 tensor.

    Converts scalar tensors (shape []) to rank-1 tensors with a single element (shape [1]).
    """
    # Scalars have a shape of lengths zero
    if len(model.get_tensor_shape(name)) == 0:
        # Lift shape to rank-1 tensor with single element
        model.set_tensor_shape(name, [1])
        # Check whether this tensor has an initializer
        if (tensor := model.get_initializer(name)) is not None:
            # Set new initializer tensor of shape [1]
            model.set_initializer(name, tensor.reshape(1))


class InferElementwiseBinaryOperation(Transformation):
    """Convert supported elementwise binary operations to their FINN custom operation."""

    @staticmethod
    def reject_output_dequant(model: ModelWrapper, node: NodeProto):
        """Filter function to filter out the last elementwise Mul operation.

        Typically filters output de-quantization operations which should happen off-chip.
        """
        # The operator must be a Mul and have no successor nodes
        if node.op_type == "Mul" and not model.find_direct_successors(node):
            # If the output is a floating-point tensors, reject this
            if model.get_tensor_datatype(node.output[0]) == "FLOAT32":
                # Filter False rejects this node
                return False
        # Filter True accepts this node
        return True

    def __init__(self, _filter=None):
        """Initialize the transformation method with an optional filter function."""
        # Initialize the base class Transformation object
        super().__init__()
        # Register the filter function as attribute
        self._filter = _filter if _filter is not None else lambda *_: True

    def apply(self, model: ModelWrapper):  # noqa
        """Apply the transform to convert elementwise binary operations to FINN custom ops."""
        # Get the model graph out of the model wrapper object
        graph = model.graph
        # Keep track of whether the graph has been modified
        graph_modified = False
        # Iterate all nodes in the graph keeping track of the index
        for index, node in enumerate(graph.node):
            # Skip transforming nodes rejected by the filter
            if not self._filter(model, node):
                continue
            # If a custom operation with corresponding name is implemented in
            # the module, this operator is supported for conversion
            if f"Elementwise{node.op_type}" in dir(elementwise_binary):
                in0 = node.input[0]
                in1 = node.input[1]
                # if both inputs are constant, throw an error and
                # ask user to run FoldConstants transform first
                assert (
                    model.get_initializer(in0) is None or model.get_initializer(in1) is None
                ), """Both inputs are constant,
                    please run FoldConstants from qonnx.transformation.fold_constants first."""
                result = node.output[0]

                # Need to "lift" potential scalar inputs to rank-1 tensors
                lift_to_rank1(in0, model)
                lift_to_rank1(in1, model)

                in0_shape = model.get_tensor_shape(in0)
                in1_shape = model.get_tensor_shape(in1)
                out_shape = model.get_tensor_shape(result)

                idt0 = model.get_tensor_datatype(in0)
                idt1 = model.get_tensor_datatype(in1)
                odt0 = model.get_tensor_datatype(result)

                new_node = helper.make_node(
                    f"Elementwise{node.op_type}",
                    [in0, in1],
                    [result],
                    domain="finn.custom_op.fpgadataflow",
                    backend="fpgadataflow",
                    lhs_shape=in0_shape,
                    rhs_shape=in1_shape,
                    out_shape=out_shape,
                    lhs_dtype=str(idt0),
                    rhs_dtype=str(idt1),
                    out_dtype=str(odt0),
                )
                graph.node.insert(index + 1, new_node)
                graph.node.remove(node)

                # Consider the graph to be modified, triggering exhaustive
                # re-application of this transformation
                graph_modified = True
                # Exiting here triggers type and shape inference and cleanup
                # after each transformed node. This helps QONNX to behave
                # better / more consistent in certain cases...
                break
        # Re-do shape and data type annotations after potential changes to the
        # model graph
        model = model.transform(InferShapes())
        model = model.transform(InferDataTypes())
        # Return the transformed model and indicate whether the graph actually
        # has been transformed
        return model, graph_modified


class InferSqueeze(Transformation):
    """Converts the Squeeze operation to the corresponding FINN custom operation."""

    # Applies the transform to a whole model graph
    def apply(self, model: ModelWrapper):  # noqa
        """Apply the transform to convert Squeeze operations to FINN custom ops."""
        # Get the model graph out of the model wrapper object
        graph = model.graph
        # Keep track of whether the graph has been modified
        graph_modified = False
        # Iterate all nodes in the graph keeping track of the index
        for index, node in enumerate(graph.node):
            # Handles Squeeze ONNX operations
            if node.op_type == "Squeeze":
                # Skip already converted nodes
                if node.domain == "finn.custom_op.fpgadataflow":
                    # Skip without warning
                    continue
                # Transplant this operator into our FINN domain
                node.domain = "finn.custom_op.fpgadataflow"  # noqa: Duplicate
                # Now we can get the CustomOp wrapper instance providing easier
                # attribute access
                inst: HWCustomOp = getCustomOp(node)
                # Set the backend attribute to mark this an operation supported
                # to be implemented on an FPGA by FINN
                inst.set_nodeattr("backend", "fpgadataflow")
                # Ge the input and output tensor names
                inp, out = node.input[0], node.output[0]
                # Set input/output shape and datatype node attributes required
                # by FINN custom op
                inst.set_nodeattr("inp_dtype", str(model.get_tensor_datatype(inp)))
                inst.set_nodeattr("inp_shape", model.get_tensor_shape(inp))
                inst.set_nodeattr("out_dtype", str(model.get_tensor_datatype(out)))
                inst.set_nodeattr("out_shape", model.get_tensor_shape(out))
                if len(node.input) > 1:
                    axes = model.get_initializer(node.input[1])
                    if np.ndim(axes) == 0:
                        # Fix axes input initializer by converting from scalar (0D) to 1D array
                        axes = np.array([axes])
                        model.set_initializer(node.input[1], axes)
                        model.set_tensor_shape(node.input[1], axes.shape)
                    # Set axes attribute (used by older opsets) even if axes is provided as input
                    inst.set_nodeattr("axes", list(axes))
                # Consider the graph to be modified, triggering exhaustive
                # re-application of this transformation
                graph_modified = True
                # Exiting here triggers type and shape inference and cleanup
                # after each transformed node. This helps QONNX to behave
                # better/more consistent in certain cases...
                break
        # Re-do shape and data type annotations after potential changes to the
        # model graph
        model = model.transform(InferShapes())
        model = model.transform(InferDataTypes())
        # Return the transformed model and indicate whether the graph actually
        # has been transformed
        return model, graph_modified


class InferUnsqueeze(Transformation):
    """Convert the Unsqueeze operation to the corresponding FINN custom operation."""

    def apply(self, model: ModelWrapper):  # noqa
        """Apply the transform to convert Unsqueeze operations to FINN custom ops."""
        # Get the model graph out of the model wrapper object
        graph = model.graph
        # Keep track of whether the graph has been modified
        graph_modified = False
        # Iterate all nodes in the graph keeping track of the index
        for index, node in enumerate(graph.node):
            # Handles Squeeze ONNX operations
            if node.op_type == "Unsqueeze":
                # Skip already converted nodes  # noqa: Duplicate
                if node.domain == "finn.custom_op.fpgadataflow":
                    # Skip without warning
                    continue
                # Transplant this operator into our FINN domain
                node.domain = "finn.custom_op.fpgadataflow"
                # Now we can get the CustomOp wrapper instance providing easier
                # attribute access
                inst: HWCustomOp = getCustomOp(node)
                # Set the backend attribute to mark this an operation supported
                # to be implemented on an FPGA by FINN
                inst.set_nodeattr("backend", "fpgadataflow")
                # Ge the input and output tensor names
                inp, out = node.input[0], node.output[0]
                # Set input/output shape and datatype node attributes required
                # by FINN custom op
                inst.set_nodeattr("inp_dtype", str(model.get_tensor_datatype(inp)))
                inst.set_nodeattr("inp_shape", model.get_tensor_shape(inp))
                inst.set_nodeattr("out_dtype", str(model.get_tensor_datatype(out)))
                inst.set_nodeattr("out_shape", model.get_tensor_shape(out))
                if len(node.input) > 1:
                    axes = model.get_initializer(node.input[1])
                    if np.ndim(axes) == 0:
                        # Fix axes input initializer by converting from scalar (0D) to 1D array
                        axes = np.array([axes])
                        model.set_initializer(node.input[1], axes)
                        model.set_tensor_shape(node.input[1], axes.shape)
                    # Set axes attribute (used by older opsets) even if axes is provided as input
                    inst.set_nodeattr("axes", list(axes))
                # Consider the graph to be modified, triggering exhaustive
                # re-application of this transformation
                graph_modified = True
                # Exiting here triggers type and shape inference and cleanup
                # after each transformed node. This helps QONNX to behave
                # better/more consistent in certain cases...
                break
        # Re-do shape and data type annotations after potential changes to the
        # model graph
        model = model.transform(InferShapes())
        model = model.transform(InferDataTypes())
        # Return the transformed model and indicate whether the graph actually
        # has been transformed
        return model, graph_modified


class InferReshape(Transformation):
    """Converts ONNX Reshape operator to the corresponding HWCustomOp."""

    def apply(self, model: ModelWrapper) -> tuple[ModelWrapper, bool]:
        """Apply the transform to convert Reshape operations hardware."""
        # Get the model graph out of the model wrapper object
        graph = model.graph
        # Keep track of whether the graph has been modified
        graph_modified = False

        # Transplant each Reshape operator into the custom domain
        for index, node in enumerate(graph.node):
            if node.op_type == "Reshape":
                # Skip already converted nodes
                if node.domain == "finn.custom_op.fpgadataflow":
                    continue

                # Skip if this is not a static-shape reshape operation
                if (out_shape := model.get_initializer(node.input[1])) is None:
                    continue

                # Skip if any of the axes is set to 0 to avoid handling
                # differences between allowzero=0 and allowzero=1...
                # TODO: This could be supported by implementing simple shape
                #  inference logic here or by normalizing shapes via streamline
                if np.any(out_shape == 0):
                    continue

                # Skip if the input shape is not statically known
                if (inp_shape := model.get_tensor_shape(node.input[0])) is None:
                    continue

                # Skip if there are more than one axes set to -1
                if list(out_shape).count(-1) > 1:
                    continue

                # The shape is allowed to contain at most one -1, in which case
                # the value is inferred from the size of the tensor and the
                # remaining dimensions
                if list(out_shape).count(-1) == 1:
                    # The following expression assumes target to be a list which
                    # allows item assignment (might be a tuple depending on the
                    # origin of target)
                    out_shape = list(out_shape)
                    # Replace the inferred dimension by the number of elements
                    # missing from the target dimension compared to the input
                    # dimension
                    out_shape[out_shape.index(-1)] = int(
                        np.prod(inp_shape) / np.abs(np.prod(out_shape))
                    )

                # Transplant from the standard ONNX domain into the FINN domain
                node.domain = "finn.custom_op.fpgadataflow"
                # Now we can get the CustomOp wrapper instance providing easier
                # attribute access
                reshape: HWCustomOp = getCustomOp(node)
                # Set the backend attribute to mark this an operation supported
                # to be implemented on an FPGA by FINN
                reshape.set_nodeattr("backend", "fpgadataflow")

                # Remove the old allowzero attribute which is not handled by the
                # hardware operator
                remove_by_name(node.attribute, "allowzero")

                # The reshape hardware operator statically annotates both, the
                # input and output shape, as node attributes
                reshape.set_nodeattr("inp_shape", list(inp_shape))
                reshape.set_nodeattr("out_shape", list(out_shape))

                # The input and output type of the reshaped tensor is the same
                # and needs to be statically annotated as an attribute
                reshape.set_nodeattr("dtype", model.get_tensor_datatype(node.input[0]).name)

                # Finally, there is no second shape input to the hardware
                # operator, thus delete this from the node input list
                node.input.pop(1)

        # Re-do shape and data type annotations after potential changes to the
        # model graph
        model = model.transform(InferShapes())
        model = model.transform(InferDataTypes())

        # Return the transformed model and indicate whether the graph actually
        # has been transformed
        return model, graph_modified
