from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from onnx import checker, helper, ModelProto, TensorProto, GraphProto, NodeProto
from typing import Sequence, Text, Tuple, List, Callable
from onnx import numpy_helper

import numpy as np  # type: ignore

import onnx.optimizer
import unittest


class TestOptimizer(unittest.TestCase):

    def _optimized(self, graph, opts):  # type: (GraphProto, Sequence[Text]) -> ModelProto
        orig_model = helper.make_model(graph, producer_name='onnx-test')
        optimized_model = onnx.optimizer.optimize(orig_model, opts)
        checker.check_model(optimized_model)
        return optimized_model

    # input_types and output_types are lists of triples of (name, type, shape)
    def _make_fake_loop_op(self,
                           body_nodes,  # type: Sequence[NodeProto]
                           input_types,  # type: Sequence[Tuple[TensorProto.DataType, Sequence[int], Text]]
                           output_types  # type: Sequence[Tuple[TensorProto.DataType, Sequence[int], Text]]
                           ):  # type: (...) -> List[NodeProto]
        zero = helper.make_tensor("trip_count_value", TensorProto.INT32, (), [10])
        true = helper.make_tensor("condition", TensorProto.BOOL, (), [True])
        # lcd is a dummy loop-carried dependency that only exists because
        # right now the schema checker is broken and assumes a variadic
        # input needs at least one value.
        graph_inputs = [helper.make_tensor_value_info("i", TensorProto.INT32, ()),
                        helper.make_tensor_value_info("cond", TensorProto.BOOL, ())]
        for type, shape, name in input_types:
            graph_inputs.append(helper.make_tensor_value_info("_" + name, type, shape))
        graph_outputs = [helper.make_tensor_value_info("cond", TensorProto.BOOL, ())]
        for type, shape, name in output_types:
            graph_outputs.append(helper.make_tensor_value_info("_" + name, type, shape))
        body_graph = helper.make_graph(body_nodes, "body_graph", graph_inputs,
                                       graph_outputs)
        loop_inputs = ["trip_count", "condition"]
        loop_inputs.extend([name for _, _, name in input_types])
        # TODO: fix checker to accept 0-input variadic inputs
        if len(loop_inputs) == 2:
            loop_inputs.append("")
        loop_outputs = [name for _, _, name in output_types]
        retval_nodes = [
            helper.make_node("Constant", [], ["trip_count"], value=zero),
            helper.make_node("Constant", [], ["condition"], value=true),
            helper.make_node("Loop", loop_inputs, loop_outputs, body=body_graph)
        ]
        return retval_nodes

    def _make_fake_if_op(self,
                         true_nodes,  # type: Sequence[NodeProto]
                         false_nodes,  # type: Sequence[NodeProto]
                         output_types  # type: Sequence[Tuple[TensorProto.DataType, Sequence[int], Text]]
                         ):  # type: (...) -> List[NodeProto]
        true = helper.make_tensor("condition", TensorProto.BOOL, (), [True])
        true_graph = helper.make_graph(true_nodes, "true_graph", [], [])
        false_graph = helper.make_graph(false_nodes, "false_graph", [], [])
        if_inputs = ["condition"]
        if_outputs = [name for _, _, name in output_types]
        retval_nodes = [
            helper.make_node("Constant", [], ["condition"], value=true),
            helper.make_node("If", if_inputs, if_outputs, then_branch=true_graph,
                             else_branch=false_graph)
        ]
        return retval_nodes

    # fn is a function that takes a single node as argument
    def _visit_all_nodes_recursive(self, graph, fn):  # type: (GraphProto, Callable[[NodeProto], None]) -> None
        for node in graph.node:
            fn(node)
            for attr in node.attribute:
                if attr.g is not None:
                    self._visit_all_nodes_recursive(attr.g, fn)
                if len(attr.graphs):
                    for gr in attr.graphs:
                        self._visit_all_nodes_recursive(gr, fn)

    def test_get_available_passes(self):  # type: () -> None
        # FIXME does not garantees to be listing all
        graph = helper.make_graph([], "dummy_graph", [], [])
        list_of_passes = onnx.optimizer.get_available_passes()
        assert isinstance(list_of_passes, (list)) and len(list_of_passes) > 0
        for pass_name in list_of_passes:
            # If pass_name is invalid it throws a RuntimeError
            self._optimized(graph, [pass_name])

    def test_eliminate_identity_single_use(self):  # type: () -> None
        nodes = [helper.make_node("Identity", ["X"], ["Y"])]
        nodes.extend(self._make_fake_loop_op(
            [helper.make_node("Identity", ["_Y"], ["_Y2"])],
            [(TensorProto.FLOAT, (5,), "Y")],
            [(TensorProto.FLOAT, (5,), "Y2")]))
        graph = helper.make_graph(
            nodes,
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (5,))],
            [helper.make_tensor_value_info("Y", TensorProto.FLOAT, (5,)),
             helper.make_tensor_value_info("Y2", TensorProto.FLOAT, (5,))])
        optimized_model = self._optimized(graph, ["eliminate_identity"])

        # All identity nodes should have been eliminated
        def check_identity(node):  # type: (NodeProto) -> None
            assert node.op_type != "Identity"
        self._visit_all_nodes_recursive(optimized_model.graph, check_identity)
        # Use of the output from the Identity node in the main graph should
        # have been replaced with the input to the identity node
        assert len(optimized_model.graph.output) == 2
        assert optimized_model.graph.output[0].name == "X"
        # Use of the output from the Identity node in the loop graph should
        # have been replaced with the input to that identity node
        assert len(optimized_model.graph.node[2].attribute[0].g.output) == 2
        assert optimized_model.graph.node[2].attribute[0].g.output[1].name == "_Y"

    def test_eliminate_identity_multiple_uses(self):  # type: () -> None
        identity = helper.make_node("Identity", ["X"], ["Y"])
        add = helper.make_node("Add", ["Z", "Y"], ["A"])
        mul = helper.make_node("Mul", ["A", "Y"], ["B"])
        graph = helper.make_graph(
            [identity, add, mul],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (5,)),
             helper.make_tensor_value_info("Z", TensorProto.FLOAT, (5,))],
            [helper.make_tensor_value_info("B", TensorProto.FLOAT, (5,))])
        optimized_model = self._optimized(graph, ["eliminate_identity"])

        for node in optimized_model.graph.node:
            assert node.op_type != "Identity"
        assert len(optimized_model.graph.node) == 2

    def test_nop_transpose(self):  # type: () -> None
        nodes = [helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 1])]
        nodes.extend(self._make_fake_loop_op(
            [helper.make_node("Transpose", ["_Y"], ["_Y2"], perm=[0, 1])],
            [(TensorProto.FLOAT, (2, 3), "Y")],
            [(TensorProto.FLOAT, (2, 3), "Y2")]))
        graph = helper.make_graph(
            nodes,
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3))],
            [helper.make_tensor_value_info("Y", TensorProto.FLOAT, (2, 3)),
             helper.make_tensor_value_info("Y2", TensorProto.FLOAT, (2, 3))])
        optimized_model = self._optimized(graph, ["eliminate_nop_transpose"])

        def check_transpose(node):  # type: (NodeProto) -> None
            assert node.op_type != "Transpose"
        self._visit_all_nodes_recursive(optimized_model.graph, check_transpose)
        # Use of the output from the Transpose node in the main graph should
        # have been replaced with the input to the identity node
        assert len(optimized_model.graph.output) == 2
        assert optimized_model.graph.output[0].name == "X"
        # Use of the output from the Transpose node in the loop graph should
        # have been replaced with the input to that identity node
        assert len(optimized_model.graph.node[2].attribute[0].g.output) == 2
        assert optimized_model.graph.node[2].attribute[0].g.output[1].name == "_Y"

    def test_nop_transpose_default(self):  # type: () -> None
        trans = helper.make_node("Transpose", ["X"], ["Y"])
        graph = helper.make_graph(
            [trans],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3))],
            [helper.make_tensor_value_info("Y", TensorProto.FLOAT, (3, 2))])
        optimized_model = self._optimized(graph, ["eliminate_nop_transpose"])

        assert len(list(optimized_model.graph.node)) == 1
        assert optimized_model.graph.node[0].op_type == "Transpose"

    def test_eliminate_unused_initializer(self):  # type: () -> None
        add = helper.make_node("Add", ["X", "Y"], ["Z"])
        graph = helper.make_graph(
            [add],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 2)),
             helper.make_tensor_value_info("Y", TensorProto.FLOAT, (1, 2))],
            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (1, 2))],
            [helper.make_tensor("A", TensorProto.FLOAT,
                                dims=(2, 3),
                                vals=np.random.randn(2, 3).astype(np.float32).tobytes(),
                                raw=True)])
        optimized_model = self._optimized(graph, ["eliminate_unused_initializer"])

        assert len(list(optimized_model.graph.initializer)) == 0

    def test_eliminate_unused_initializer_input(self):  # type: () -> None
        add = helper.make_node("Add", ["X", "Y"], ["Z"])
        graph = helper.make_graph(
            [add],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 2)),
             helper.make_tensor_value_info("Y", TensorProto.FLOAT, (1, 2)),
             helper.make_tensor_value_info("A", TensorProto.FLOAT, (2, 3))],
            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (1, 2))],
            [helper.make_tensor("A", TensorProto.FLOAT,
                                dims=(2, 3),
                                vals=np.random.randn(2, 3).astype(np.float32).tobytes(),
                                raw=True)])
        optimized_model = self._optimized(graph, ["eliminate_unused_initializer"])

        assert len(list(optimized_model.graph.initializer)) == 0
        assert len(optimized_model.graph.input) == 2

    def test_eliminate_unused_initializer_no_eliminate_used_default(self):  # type: () -> None
        add = helper.make_node("Add", ["X", "A"], ["Z"])
        graph = helper.make_graph(
            [add],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 2)),
             helper.make_tensor_value_info("A", TensorProto.FLOAT, (1, 2))],
            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (1, 2))],
            [helper.make_tensor("A", TensorProto.FLOAT,
                                dims=(1, 2),
                                vals=np.random.randn(1, 2).astype(np.float32).tobytes(),
                                raw=True)])
        optimized_model = self._optimized(graph, ["eliminate_unused_initializer"])

        assert len(list(optimized_model.graph.initializer)) == 1

    def test_eliminate_unused_initializer_no_eliminate_used(self):  # type: () -> None
        nodes = [helper.make_node("Add", ["X", "A"], ["Z"])]
        nodes.extend(self._make_fake_loop_op(
            [helper.make_node("Add", ["_X", "_A"], ["_Z2"])],
            [(TensorProto.FLOAT, (1, 2), "X"),
             (TensorProto.FLOAT, (1, 2), "A")],
            [(TensorProto.FLOAT, (1, 2), "Z2")]))
        graph = helper.make_graph(
            nodes,
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 2)),
             helper.make_tensor_value_info("A", TensorProto.FLOAT, (1, 2))],
            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (1, 2))],
            [helper.make_tensor("A", TensorProto.FLOAT,
                                dims=(1, 2),
                                vals=np.random.randn(1, 2).astype(np.float32).tobytes(),
                                raw=True)])
        optimized_model = self._optimized(graph, ["eliminate_unused_initializer"])

        # Add, Constant (trip count), Constant (cond), Loop
        assert len(list(optimized_model.graph.node)) == 4
        assert optimized_model.graph.node[0].op_type == "Add"
        assert optimized_model.graph.output[0].name == "Z"
        # Add
        assert len(optimized_model.graph.node[3].attribute[0].g.node) == 1
        assert optimized_model.graph.node[3].attribute[0].g.node[0].op_type == 'Add'
        assert optimized_model.graph.node[3].attribute[0].g.output[1].name == '_Z2'

        assert len(list(optimized_model.graph.initializer)) == 1

    def test_eliminate_unused_initializer_no_eliminate_output(self):  # type: () -> None
        add = helper.make_node("Add", ["X", "Y"], ["Z"])
        graph = helper.make_graph(
            [add],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 2)),
             helper.make_tensor_value_info("Y", TensorProto.FLOAT, (1, 2)),
             helper.make_tensor_value_info("A", TensorProto.FLOAT, (2, 3))],
            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (1, 2)),
             helper.make_tensor_value_info("A", TensorProto.FLOAT, (2, 3))],
            [helper.make_tensor("A", TensorProto.FLOAT,
                                dims=(2, 3),
                                vals=np.random.randn(2, 3).astype(np.float32).tobytes(),
                                raw=True)])
        optimized_model = self._optimized(graph, ["eliminate_unused_initializer"])

        assert len(list(optimized_model.graph.initializer)) == 1
        assert "Z" in [o.name for o in optimized_model.graph.output]

    def test_extract_constant_to_initializer(self):  # type: () -> None
        conv = helper.make_node("Conv", ["X", "Y"], ["Z"])
        constant = helper.make_node("Constant", [], ["A"],
                                    value=helper.make_tensor(
                                        name="bias",
                                        data_type=TensorProto.FLOAT,
                                        dims=(16, 1, 1),
                                        vals=np.random.randn(16).astype(np.float32).tolist()))
        add = helper.make_node("Add", ["Z", "A"], ["B"])
        graph = helper.make_graph(
            [conv, constant, add],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 5, 3, 3)),
             helper.make_tensor_value_info("Y", TensorProto.FLOAT, (16, 5, 3, 3))],
            [helper.make_tensor_value_info("B", TensorProto.FLOAT, (1, 16, 1, 1))],
        )
        optimized_model = self._optimized(graph, ["extract_constant_to_initializer"])
        self.assertEqual(
            set(vi.name for vi in optimized_model.graph.input),
            {'X', 'Y', 'A'})

        self.assertEqual(len(optimized_model.graph.initializer), 1)
        init = optimized_model.graph.initializer[0]
        self.assertEqual(init.name, 'A')
        self.assertEqual(init.dims, [16, 1, 1])
        self.assertEqual(init.data_type, TensorProto.FLOAT)

        self.assertEqual([n.op_type for n in optimized_model.graph.node], ['Conv', 'Add'])

    def test_fuse_transpose(self):  # type: () -> None
        nodes = [helper.make_node("Transpose", ["X"], ["Y"], perm=[1, 0, 2]),
                 helper.make_node("Transpose", ["Y"], ["Z"], perm=[2, 0, 1]),
                 helper.make_node("Transpose", ["Z"], ["A"], perm=[2, 0, 1])]
        nodes.extend(self._make_fake_loop_op(
            [helper.make_node("Transpose", ["_X"], ["_Y2"], perm=[1, 0, 2]),
             helper.make_node("Transpose", ["_Y2"], ["_Y3"], perm=[2, 0, 1]),
             helper.make_node("Transpose", ["_Y3"], ["_Y4"], perm=[2, 0, 1])],
            [(TensorProto.FLOAT, (2, 3), "X")],
            [(TensorProto.FLOAT, (2, 3), "Y4")]))
        graph = helper.make_graph(
            nodes,
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3, 4))],
            [helper.make_tensor_value_info("A", TensorProto.FLOAT, (4, 3, 2)),
             helper.make_tensor_value_info("Y4", TensorProto.FLOAT, (4, 3, 2))])
        optimized_model = self._optimized(graph, ["fuse_consecutive_transposes"])

        # Transpose, Constant (trip count), Constant (cond), Loop
        assert len(list(optimized_model.graph.node)) == 4
        # Transpose
        assert len(optimized_model.graph.node[3].attribute[0].g.node) == 1

    def test_fuse_transpose_default(self):  # type: () -> None
        trans1 = helper.make_node("Transpose", ["X"], ["Y"])
        trans2 = helper.make_node("Transpose", ["Y"], ["Z"])
        graph = helper.make_graph(
            [trans1, trans2],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3, 4))],
            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (2, 3, 4))])
        optimized_model = self._optimized(graph, ["fuse_consecutive_transposes"])

        assert len(list(optimized_model.graph.node)) == 0

    def test_fuse_transpose_default_no_fuse(self):  # type: () -> None
        trans1 = helper.make_node("Transpose", ["X"], ["Y"])
        trans2 = helper.make_node("Transpose", ["Y"], ["Z"], perm=[0, 1, 2])
        graph = helper.make_graph(
            [trans1, trans2],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3, 4))],
            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (4, 3, 2))])
        optimized_model = self._optimized(graph, ["fuse_consecutive_transposes"])

        assert len(list(optimized_model.graph.node)) == 2
        for node in optimized_model.graph.node:
            assert node.op_type == "Transpose"

    def test_fuse_transpose_into_gemm(self):  # type: () -> None
        nodes = [helper.make_node("Transpose", ["X"], ["A"], perm=[1, 0]),
                 helper.make_node("Transpose", ["Y"], ["B"], perm=[1, 0]),
                 helper.make_node("Gemm", ["A", "B", "C"], ["Z"])]
        nodes.extend(self._make_fake_loop_op(
            [helper.make_node("Transpose", ["_X"], ["_A"], perm=[1, 0]),
             helper.make_node("Transpose", ["_Y"], ["_B"], perm=[1, 0]),
             helper.make_node("Gemm", ["_A", "_B", "_C"], ["_Z2"])],
            [(TensorProto.FLOAT, (2, 3), "X"),
             (TensorProto.FLOAT, (5, 2), "Y"),
             (TensorProto.FLOAT, (3, 5), "C")],
            [(TensorProto.FLOAT, (2, 3), "Z2")]))
        graph = helper.make_graph(
            nodes,
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3)),
             helper.make_tensor_value_info("Y", TensorProto.FLOAT, (5, 2)),
             helper.make_tensor_value_info("C", TensorProto.FLOAT, (3, 5))],
            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (3, 5))])
        optimized_model = self._optimized(graph, ["fuse_transpose_into_gemm"])

        # Gemm, Constant (trip count), Constant (cond), Loop
        assert len(list(optimized_model.graph.node)) == 4
        assert optimized_model.graph.node[0].op_type == "Gemm"
        # Gemm
        assert len(optimized_model.graph.node[3].attribute[0].g.node) == 1
        assert optimized_model.graph.node[3].attribute[0].g.node[0].op_type == "Gemm"

    def test_fuse_add_bias_into_conv_use_weight_shape(self):  # type: () -> None
        nodes = [helper.make_node("Conv", ["X", "Y"], ["Z"]),
                 helper.make_node("Add", ["Z", "A"], ["B"])]
        nodes.extend(self._make_fake_loop_op(
            [helper.make_node("Conv", ["_X", "_Y"], ["_Z"]),
             helper.make_node("Add", ["_Z", "_A"], ["_B2"])],
            [(TensorProto.FLOAT, (1, 5, 3, 3), "X"),
             (TensorProto.FLOAT, (16, 5, 3, 3), "Y"),
             (TensorProto.FLOAT, (16, 1, 1), "A")],
            [(TensorProto.FLOAT, (1, 16, 3, 3), "B2")]))
        graph = helper.make_graph(
            nodes,
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 5, 3, 3)),
             helper.make_tensor_value_info("Y", TensorProto.FLOAT, (16, 5, 3, 3)),
             helper.make_tensor_value_info("A", TensorProto.FLOAT, (16, 1, 1))],
            [helper.make_tensor_value_info("B", TensorProto.FLOAT, (1, 16, 1, 1))],
        )
        optimized_model = self._optimized(graph, ["fuse_add_bias_into_conv"])

        # Squeeze, Conv, Constant (trip count), Constant (condition), Loop
        assert len(list(optimized_model.graph.node)) == 5
        assert optimized_model.graph.node[0].op_type == 'Squeeze'
        assert optimized_model.graph.node[1].op_type == 'Conv'
        assert optimized_model.graph.output[0].name == 'Z'
        # Squeeze, Conv
        assert len(optimized_model.graph.node[4].attribute[0].g.node) == 2
        assert optimized_model.graph.node[4].attribute[0].g.node[0].op_type == 'Squeeze'
        assert optimized_model.graph.node[4].attribute[0].g.node[1].op_type == 'Conv'
        # Output 1 since 0 is 'cond'
        assert optimized_model.graph.node[4].attribute[0].g.output[1].name == '_Z'

    def test_fuse_add_bias_into_conv_use_weight_shape_with_tile(self):  # type: () -> None
        conv = helper.make_node("Conv", ["X", "Y"], ["Z"])
        add = helper.make_node("Add", ["Z", "A"], ["B"])
        graph = helper.make_graph(
            [conv, add],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 5, 3, 3)),
             helper.make_tensor_value_info("Y", TensorProto.FLOAT, (16, 5, 3, 3)),
             helper.make_tensor_value_info("A", TensorProto.FLOAT, (1,))],
            [helper.make_tensor_value_info("B", TensorProto.FLOAT, (1, 16, 1, 1))],
        )
        optimized_model = self._optimized(graph, ["fuse_add_bias_into_conv"])

        assert len(list(optimized_model.graph.node)) == 3
        assert len(optimized_model.graph.value_info) == 1
        assert optimized_model.graph.value_info[0].type.tensor_type.elem_type == TensorProto.INT64
        assert len(optimized_model.graph.value_info[0].type.tensor_type.shape.dim) == 1
        assert optimized_model.graph.node[0].op_type == 'Constant'
        assert optimized_model.graph.node[1].op_type == 'Tile'
        assert optimized_model.graph.node[2].op_type == 'Conv'
        assert optimized_model.graph.output[0].name == 'Z'

    def test_fuse_add_bias_into_conv_use_conv_shape(self):  # type: () -> None
        sub = helper.make_node("Sub", ["M", "N"], ["Y"])
        conv = helper.make_node("Conv", ["X", "Y"], ["Z"])
        add = helper.make_node("Add", ["Z", "A"], ["B"])
        graph = helper.make_graph(
            [sub, conv, add],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 5, 3, 3)),
             helper.make_tensor_value_info("M", TensorProto.FLOAT, (16, 5, 3, 3)),
             helper.make_tensor_value_info("N", TensorProto.FLOAT, (16, 5, 3, 3)),
             helper.make_tensor_value_info("A", TensorProto.FLOAT, (1, 16, 1, 1))],
            [helper.make_tensor_value_info("B", TensorProto.FLOAT, (1, 16, 1, 1))],
            value_info=[
                helper.make_tensor_value_info("Z", TensorProto.FLOAT, (1, 16, 1, 1))
            ],
        )
        optimized_model = self._optimized(graph, ["fuse_add_bias_into_conv"])

        assert len(optimized_model.graph.node) == 3
        assert optimized_model.graph.node[0].op_type == 'Sub'
        assert optimized_model.graph.node[1].op_type == 'Squeeze'
        assert optimized_model.graph.node[2].op_type == 'Conv'
        assert optimized_model.graph.output[0].name == 'Z'
        assert optimized_model.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT
        assert len(optimized_model.graph.output[0].type.tensor_type.shape.dim) == 4

    def test_fuse_add_bias_into_conv_use_move_constant(self):  # type: () -> None
        conv = helper.make_node("Conv", ["X", "Y"], ["Z"])
        constant = helper.make_node("Constant", [], ["A"],
                                    value=helper.make_tensor(
                                        name="bias",
                                        data_type=TensorProto.FLOAT,
                                        dims=(16, 1, 1),
                                        vals=np.random.randn(16).astype(np.float32).tolist()))
        add = helper.make_node("Add", ["Z", "A"], ["B"])
        graph = helper.make_graph(
            [conv, constant, add],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 5, 3, 3)),
             helper.make_tensor_value_info("Y", TensorProto.FLOAT, (16, 5, 3, 3))],
            [helper.make_tensor_value_info("B", TensorProto.FLOAT, (1, 16, 1, 1))],
            value_info=[
                helper.make_tensor_value_info("A", TensorProto.FLOAT, (16, 1, 1)),
            ]
        )
        optimized_model = self._optimized(graph, ["fuse_add_bias_into_conv"])

        assert len(optimized_model.graph.node) == 3
        assert optimized_model.graph.node[0].op_type == 'Constant'
        assert optimized_model.graph.node[1].op_type == 'Squeeze'
        assert optimized_model.graph.node[2].op_type == 'Conv'
        assert optimized_model.graph.output[0].name == 'Z'
        assert optimized_model.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT
        assert len(optimized_model.graph.output[0].type.tensor_type.shape.dim) == 4

    def test_fuse_add_bias_into_conv_squeeze_1d_bias_no_fuse(self):  # type: () -> None
        conv = helper.make_node("Conv", ["X", "Y"], ["Z"])
        add = helper.make_node("Add", ["Z", "A"], ["B"])
        graph = helper.make_graph(
            [conv, add],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 5, 3, 3)),
             helper.make_tensor_value_info("Y", TensorProto.FLOAT, (16, 5, 3, 3)),
             helper.make_tensor_value_info("A", TensorProto.FLOAT, (3,))],
            [helper.make_tensor_value_info("B", TensorProto.FLOAT, (1, 16, 1, 3))],
            value_info=[
                helper.make_tensor_value_info("Z", TensorProto.FLOAT, (1, 16, 1, 1)),
            ]
        )
        optimized_model = self._optimized(graph, ["fuse_add_bias_into_conv"])

        assert len(list(optimized_model.graph.node)) == 2
        assert optimized_model.graph.node[0].op_type == 'Conv'
        assert optimized_model.graph.node[1].op_type == 'Add'

    def test_fuse_add_bias_into_conv_squeeze_3d_bias_no_fuse(self):  # type: () -> None
        conv = helper.make_node("Conv", ["X", "Y"], ["Z"])
        add = helper.make_node("Add", ["Z", "A"], ["B"])
        graph = helper.make_graph(
            [conv, add],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 5, 3, 3)),
             helper.make_tensor_value_info("Y", TensorProto.FLOAT, (16, 5, 3, 3)),
             helper.make_tensor_value_info("A", TensorProto.FLOAT, (16, 3, 3))],
            [helper.make_tensor_value_info("B", TensorProto.FLOAT, (1, 16, 3, 3))],
            value_info=[
                helper.make_tensor_value_info("Z", TensorProto.FLOAT, (1, 16, 1, 1)),
            ]
        )
        optimized_model = self._optimized(graph, ["fuse_add_bias_into_conv"])

        assert len(list(optimized_model.graph.node)) == 2
        assert optimized_model.graph.node[0].op_type == 'Conv'
        assert optimized_model.graph.node[1].op_type == 'Add'

    def test_fuse_add_bias_into_conv_squeeze_4d_bias_no_fuse(self):  # type: () -> None
        conv = helper.make_node("Conv", ["X", "Y"], ["Z"])
        add = helper.make_node("Add", ["Z", "A"], ["B"])
        graph = helper.make_graph(
            [conv, add],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 5, 3, 3)),
             helper.make_tensor_value_info("Y", TensorProto.FLOAT, (16, 5, 3, 3)),
             helper.make_tensor_value_info("A", TensorProto.FLOAT, (1, 16, 3, 3))],
            [helper.make_tensor_value_info("B", TensorProto.FLOAT, (1, 16, 3, 3))]
        )
        optimized_model = self._optimized(graph, ["fuse_add_bias_into_conv"])

        assert len(list(optimized_model.graph.node)) == 2
        assert optimized_model.graph.node[0].op_type == 'Conv'
        assert optimized_model.graph.node[1].op_type == 'Add'

    def test_fuse_consecutive_squeezes(self):  # type: () -> None
        nodes = [helper.make_node("Squeeze", ["X"], ["Y"], axes=[0, 4, 5]),
                 helper.make_node("Squeeze", ["Y"], ["Z"], axes=[0, 3])]
        nodes.extend(self._make_fake_loop_op(
            [helper.make_node("Squeeze", ["_X"], ["_Y"], axes=[0, 4, 5]),
             helper.make_node("Squeeze", ["_Y"], ["_Z2"], axes=[0, 3])],
            [(TensorProto.FLOAT, (1, 1, 2, 3, 1, 1, 1, 1, 8, 9), "X")],
            [(TensorProto.FLOAT, (2, 3, 1, 8, 9), "Z2")]))

        graph = helper.make_graph(
            nodes,
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 1, 2, 3, 1, 1, 1, 1, 8, 9))],
            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (2, 3, 1, 8, 9))])
        optimized_model = self._optimized(graph, ["fuse_consecutive_squeezes"])

        # Squeeze, Constant (trip count), Constant (cond), Loop
        assert optimized_model.graph.node[0].op_type == "Squeeze"
        assert list(optimized_model.graph.node[0].attribute[0].ints) == [0, 1, 4, 5, 6]
        assert len(list(optimized_model.graph.node)) == 4

    def test_fuse_consecutive_squeezes_default(self):  # type: () -> None
        squeeze1 = helper.make_node("Squeeze", ["X"], ["Y"], axes=[0, 4, 5])
        squeeze2 = helper.make_node("Squeeze", ["Y"], ["Z"], axes=[0, 3])
        squeeze3 = helper.make_node("Squeeze", ["Z"], ["A"], axes=[2])
        nodes = [squeeze1, squeeze2, squeeze3]
        graph = helper.make_graph(
            nodes,
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 1, 2, 3, 1, 1, 1, 1, 8, 9))],
            [helper.make_tensor_value_info("A", TensorProto.FLOAT, (2, 3, 8, 9))])
        optimized_model = self._optimized(graph, ["fuse_consecutive_squeezes"])

        assert optimized_model.graph.node[0].op_type == "Squeeze"
        assert list(optimized_model.graph.node[0].attribute[0].ints) == [0, 1, 4, 5, 6, 7]
        assert len(list(optimized_model.graph.node)) == 1

    def test_fuse_consecutive_squeezes_random(self):  # type: () -> None
        x_shape = [1, 1, 1, 3, 4, 1, 6, 1, 1, 9]
        s1_one_indices = [i for i, a in enumerate(x_shape) if a == 1]
        s1_axes = np.random.choice(s1_one_indices, size=np.random.randint(low=1, high=len(s1_one_indices) - 1),
                                   replace=False)
        s2_x_shape = [a for i, a in enumerate(x_shape) if i not in s1_axes]
        s2_one_indices = [i for i, a in enumerate(s2_x_shape) if a == 1]
        s2_axes = s2_one_indices

        squeeze1 = helper.make_node("Squeeze", ["X"], ["Y"], axes=s1_axes)
        squeeze2 = helper.make_node("Squeeze", ["Y"], ["Z"], axes=s2_axes)
        nodes = [squeeze1, squeeze2]
        graph = helper.make_graph(
            nodes,
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, x_shape)],
            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (3, 4, 6, 9))])
        optimized_model = self._optimized(graph, ["fuse_consecutive_squeezes"])

        assert optimized_model.graph.node[0].op_type == "Squeeze"
        assert list(optimized_model.graph.node[0].attribute[0].ints) == [0, 1, 2, 5, 7, 8]
        assert len(list(optimized_model.graph.node)) == 1

    def test_fuse_consecutive_squeezes_multi_uses(self):  # type: () -> None
        squeeze1 = helper.make_node("Squeeze", ["X"], ["Y"], axes=[0, 4, 5])
        add = helper.make_node("Add", ["Y", "A"], ["Z2"])
        squeeze2 = helper.make_node("Squeeze", ["Y"], ["Z"], axes=[0, 3])
        graph = helper.make_graph(
            [squeeze1, add, squeeze2],
            "test",
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 1, 2, 3, 1, 1, 1, 1, 8, 9)),
             helper.make_tensor_value_info("A", TensorProto.FLOAT, (1,))],
            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (2, 3, 1, 8, 9)),
             helper.make_tensor_value_info("Z2", TensorProto.FLOAT, (1, 2, 3, 1, 1, 8, 9))])
        optimized_model = self._optimized(graph, ["fuse_consecutive_squeezes"])

        assert optimized_model.graph.node[0].op_type == "Squeeze"
        assert list(optimized_model.graph.node[0].attribute[0].ints) == [0, 4, 5]
        assert optimized_model.graph.node[2].op_type == "Squeeze"
        assert optimized_model.graph.node[2].input == ["X"]
        assert list(optimized_model.graph.node[2].attribute[0].ints) == [0, 1, 4, 5, 6]