onnx--onnx/onnx/backend/test/case/node/convinteger.py

# Copyright (c) ONNX Project Contributors
#
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations

import numpy as np

import onnx
from onnx.backend.test.case.base import Base
from onnx.backend.test.case.node import expect


class ConvInteger(Base):
    @staticmethod
    def export_without_padding() -> None:
        x = (
            np.array([2, 3, 4, 5, 6, 7, 8, 9, 10])
            .astype(np.uint8)
            .reshape((1, 1, 3, 3))
        )
        x_zero_point = np.uint8(1)
        w = np.array([1, 1, 1, 1]).astype(np.uint8).reshape((1, 1, 2, 2))

        y = np.array([12, 16, 24, 28]).astype(np.int32).reshape(1, 1, 2, 2)

        # ConvInteger without padding
        convinteger_node = onnx.helper.make_node(
            "ConvInteger", inputs=["x", "w", "x_zero_point"], outputs=["y"]
        )

        expect(
            convinteger_node,
            inputs=[x, w, x_zero_point],
            outputs=[y],
            name="test_convinteger_without_padding",
        )

    @staticmethod
    def export_with_padding() -> None:
        x = (
            np.array([2, 3, 4, 5, 6, 7, 8, 9, 10])
            .astype(np.uint8)
            .reshape((1, 1, 3, 3))
        )
        x_zero_point = np.uint8(1)
        w_zero_points = np.array([0, 1], dtype=np.uint8)
        w = np.array([1, 1, 1, 1, 1, 1, 1, 1]).astype(np.uint8).reshape((2, 1, 2, 2))

        y = (
            np.array(
                [
                    1,
                    3,
                    5,
                    3,
                    5,
                    12,
                    16,
                    9,
                    11,
                    24,
                    28,
                    15,
                    7,
                    15,
                    17,
                    9,
                    0,
                    0,
                    0,
                    0,
                    0,
                    0,
                    0,
                    0,
                    0,
                    0,
                    0,
                    0,
                    0,
                    0,
                    0,
                    0,
                ]
            )
            .astype(np.int32)
            .reshape((1, 2, 4, 4))
        )

        # ConvInteger with padding
        convinteger_node_with_padding = onnx.helper.make_node(
            "ConvInteger",
            inputs=["x", "w", "x_zero_point", "w_zero_points"],
            outputs=["y"],
            pads=[1, 1, 1, 1],
        )

        expect(
            convinteger_node_with_padding,
            inputs=[x, w, x_zero_point, w_zero_points],
            outputs=[y],
            name="test_convinteger_with_padding",
        )