#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import copy
from typing import List

import numpy as np
import onnx
import onnx_graphsurgeon as gs
import pytest
from onnx_graphsurgeon.ir.function import Function
from onnx_graphsurgeon.ir.graph import Graph
from onnx_graphsurgeon.ir.node import Node
from onnx_graphsurgeon.ir.tensor import Constant, LazyValues, Tensor, Variable
from onnx_graphsurgeon.logger import G_LOGGER
from onnx_graphsurgeon.util import misc
from onnx_graphsurgeon.util.exception import OnnxGraphSurgeonException
from onnx_graphsurgeon.util.misc import SynchronizedList
from onnx_models import const_foldable, shape_cast_elision

G_LOGGER.severity = G_LOGGER.ULTRA_VERBOSE


@Graph.register()
def shape(self, inp):
    return self.layer(op="Shape", inputs=[inp], outputs=["shape_out"])[0]


@Graph.register()
def cast(self, inp, to):
    return self.layer(op="Cast", inputs=[inp], outputs=["cast_out"], attrs={"to": to})[
        0
    ]


@Graph.register()
def constant(self, values):
    return self.layer(
        op="Constant",
        inputs=[],
        outputs=["constant_out"],
        attrs={"value": Constant("values", values)},
    )[0]


@Graph.register()
def identity(self, inp):
    out = self.layer(op="Identity", inputs=[inp], outputs=["identity_out"])[0]
    out.dtype = inp.dtype
    return out


@Graph.register()
def relu(self, inp):
    out = self.layer(op="Relu", inputs=[inp], outputs=["relu_out"])[0]
    out.dtype = inp.dtype
    return out


@Graph.register()
def add(self, a, b, name=None):
    outputs = [Variable(name=name)] if name else ["add_out"]
    out = self.layer(op="Add", inputs=[a, b], outputs=outputs)[0]
    out.dtype = a.dtype or b.dtype
    return out


@Graph.register()
def mul(self, a, b, name=None):
    outputs = [Variable(name=name)] if name else ["mul_out"]
    out = self.layer(op="Mul", inputs=[a, b], outputs=outputs)[0]
    out.dtype = a.dtype or b.dtype
    return out


@Graph.register()
def less(self, a, b, name=None):
    outputs = [Variable(name=name)] if name else ["less_out"]
    out = self.layer(op="Less", inputs=[a, b], outputs=outputs)[0]
    out.dtype = bool
    return out


# A fake op that can be used to ensure things work even when there is an invalid
# node present in the model.
@Graph.register()
def fake(self, inp, name=None):
    outputs = [Variable(name=name)] if name else ["fake_out"]
    out = self.layer(op="Fake", inputs=[inp], outputs=outputs)[0]
    out.dtype = inp.dtype
    return out


@gs.Graph.register()
def gather(self, data, indices):
    return self.layer(op="Gather", inputs=[data, indices], outputs=["gather_out"])[0]


@gs.Graph.register()
def slice(self, data, starts=None, ends=None, axes=None, steps=None):
    inputs = []
    for inp in [data, starts, ends, axes, steps]:
        if inp is None:
            break
        inputs.append(inp)

    return self.layer(op="Slice", inputs=inputs, outputs=["slice_out"])[0]


@gs.Graph.register()
def nested(self, inp, graph):
    return self.layer(
        op="Nested", inputs=[inp], outputs=["nested_out"], attrs={"body": graph}
    )[0]


@gs.Graph.register()
def if_op(self, cond, then_graph, else_graph):
    return self.layer(
        op="If",
        inputs=[cond],
        outputs=["if_out"],
        attrs={"then_branch": then_graph, "else_branch": else_graph},
    )[0]


@gs.Graph.register()
def tile(self, inp, repeats):
    out = self.layer(op="Tile", inputs=[inp, repeats], outputs=["tile_out"])[0]
    out.dtype = inp.dtype
    return out


@gs.Graph.register()
def dequantize_linear(self, inp, scale, zero_point, axis=1):
    out = self.layer(
        op="DequantizeLinear",
        inputs=[inp, scale, zero_point],
        outputs=["dequantize_linear_out"],
        attrs={"axis": axis},
    )[0]
    out.dtype = np.float32
    return out


@gs.Graph.register()
def quantize_linear(self, inp, out_scale, out_zero_point, axis=1):
    out = self.layer(
        op="QuantizeLinear",
        inputs=[inp, out_scale, out_zero_point],
        outputs=["quantize_linear_out"],
        attrs={"axis": axis},
    )[0]
    out.dtype = np.int8
    return out


@gs.Graph.register()
def pad(self, data, pads, constant_value=None):
    constant_value = misc.default_value(constant_value, Variable.empty())
    out = self.layer(
        op="Pad", inputs=[data, pads, constant_value], outputs=["pad_out"]
    )[0]
    out.dtype = data.dtype
    return out


@gs.Graph.register()
def softmax(self, data, axis=None):
    attrs = {}
    if axis is not None:
        attrs["axis"] = axis
    out = self.layer(op="Softmax", inputs=[data], outputs=["softmax_out"], attrs=attrs)[
        0
    ]
    out.dtype = data.dtype
    return out


# Generates a graph where an outer node has no outputs except
# within the subgraph. ONNX-GS should recognize that the node
# is being used, and should not remove it during cleanup().
def make_nested_graph():
    inp = Variable("input")
    id_out = Variable("id_out")
    identity = Node(op="Identity", inputs=[inp], outputs=[id_out])

    # Subgraph outputs come from the parent node, but nodes in the subgraph
    # can use nodes from the outer graphs too.
    subgraph_inputs = [Variable("subgraph_inp")]
    subgraph_id_out = Variable("subgraph_id_out")
    subgraph_outputs = [Variable("subgraph_out")]

    subgraph_identity0 = Node(op="Identity", inputs=[id_out], outputs=[subgraph_id_out])
    subgraph_identity1 = Node(
        op="Identity", inputs=[subgraph_id_out], outputs=subgraph_outputs
    )

    subgraph = Graph(
        nodes=[subgraph_identity0, subgraph_identity1],
        inputs=subgraph_inputs,
        outputs=subgraph_outputs,
    )

    nested_out = Variable("nested_out")
    nested_node = Node(
        op="Nested", attrs={"body": subgraph}, inputs=[inp], outputs=[nested_out]
    )

    return Graph(nodes=[identity, nested_node], inputs=[inp], outputs=[nested_out])


@pytest.fixture
def nested_graph():
    yield make_nested_graph()


@pytest.fixture
def very_nested_graph():
    inner_subgraph_1 = Graph(name="inner_subgraph_1")
    inner_subgraph_2 = Graph(name="inner_subgraph_2")
    inner_subgraph_3 = Graph(name="inner_subgraph_3")
    outer_subgraph_1 = Graph(name="subgraph1")
    outer_subgraph_2 = Graph(
        name="subgraph2",
        nodes=[Node("Add", attrs={"x": inner_subgraph_1, "y": inner_subgraph_2})],
    )
    outer_subgraph_3 = Graph(
        name="subgraph3", nodes=[Node("Add", attrs={"x": inner_subgraph_3, "y": 3.14})]
    )
    node_1 = Node(
        op="Add",
        attrs={
            "x": outer_subgraph_1,
            "y": outer_subgraph_2,
            "z": 5,
            "w": outer_subgraph_3,
        },
    )
    node_2 = Node(op="Add", attrs={"x": outer_subgraph_3})
    return Graph(nodes=[node_1, node_2], name="main_graph")


class TestBasic(object):
    def test_generate_name(self):
        graph = Graph()
        generated_names = set()
        existing_names = {
            "name_{}".format(i) for i in range(50, 150)
        }  # names_50 to names_149
        num_names = 100
        for idx in range(num_names):
            generated_names.add(graph._generate_name("name", existing_names))
        assert len(generated_names) == num_names  # 100 unique generated_names
        assert (
            len(generated_names.intersection(existing_names)) == 0
        )  # no generated_names in existing_names
        expected_names = {"name_{}".format(i) for i in range(0, 50)}
        expected_names.update({"name_{}".format(i) for i in range(150, 200)})
        assert (
            generated_names == expected_names
        )  # expect 'names_0' to 'names_49', 'names_150' to 'names_199'

    def test_equal(self, nested_graph):
        assert nested_graph == nested_graph

    def test_equal_inputs_unequal(self):
        g0 = make_nested_graph()
        g1 = make_nested_graph()

        g0.inputs.append(Variable("test"))

        assert not (g0 == g1)

    def test_equal_outputs_unequal(self):
        g0 = make_nested_graph()
        g1 = make_nested_graph()

        g0.outputs.append(Variable("test"))

        assert not (g0 == g1)

    def test_equal_nested_unequal(self):
        g0 = make_nested_graph()
        g1 = make_nested_graph()

        # Changing the nested subgraph should make the graphs unequal
        g0.nodes[1].inputs[0].name = "subgraph_inp_modified"

        assert not (g0 == g1)

    def test_subgraphs_not_recursive(self, very_nested_graph):
        unrelated_graph = Graph(name="unrelated")
        subgraph_names = {subgraph.name for subgraph in very_nested_graph.subgraphs()}
        assert subgraph_names == {"subgraph1", "subgraph2", "subgraph3"}

    def test_subgraphs_recursive(self, very_nested_graph):
        unrelated_graph = Graph(name="unrelated")
        subgraph_names = {
            subgraph.name for subgraph in very_nested_graph.subgraphs(recursive=True)
        }
        assert subgraph_names == {
            "subgraph1",
            "subgraph2",
            "subgraph3",
            "inner_subgraph_1",
            "inner_subgraph_2",
            "inner_subgraph_3",
        }


class TestRegister(object):
    def test_register(self):
        @Graph.register()
        def fake_add(self, a, b):
            return self.layer(op="Add", inputs=[a, b], outputs=["add_out"])

        graph = Graph()
        [output] = graph.fake_add("a", "b")
        assert "add_out" in output.name
        assert len(graph.nodes) == 1
        assert graph.nodes[-1].op == "Add"

    def test_register_opset(self):
        @Graph.register(opsets=[11])
        def fake_add(self, a, b):
            return self.layer(op="Add", inputs=[a, b], outputs=["add_out"])

        @Graph.register(opsets=[10])
        def fake_add(self, a, b):
            return self.layer(op="Add-10", inputs=[a, b], outputs=["add_out"])

        graph = Graph()
        [output] = graph.fake_add("a", "b")
        assert "add_out" in output.name
        assert len(graph.nodes) == 1
        assert graph.nodes[-1].op == "Add"

        graph_opset10 = Graph(opset=10)
        [output] = graph_opset10.fake_add("a", "b")
        assert "add_out" in output.name
        assert len(graph_opset10.nodes) == 1
        assert graph_opset10.nodes[-1].op == "Add-10"

    def test_register_name_conflict(self):
        @Graph.register()
        def fake_mul(self, a, b):
            return self.layer(
                op="Add", domain="domain1", inputs=[a, b], outputs=["mul_out"]
            )

        func = Function("fake_mul", domain="domain2")
        graph = Graph(functions=[func])
        graph.fake_mul("a", "b")
        assert len(graph.nodes) == 1
        assert graph.nodes[0].domain == "domain1"


class TestLayer(object):
    def test_layer_default_naming(self):
        node1 = Node(
            name="onnx_graphsurgeon_node_0", op="Identity"
        )  # injecting default name
        node2 = Node(
            name="onnx_graphsurgeon_node_1", op="Identity"
        )  # injecting default name again
        graph = Graph(nodes=[node1, node2])
        graph.layer(
            op="Identity"
        )  # new default name should be onnx_graphsurgeon_node_2
        assert graph.nodes[-1].name == "onnx_graphsurgeon_node_2"
        graph.layer(
            op="Identity"
        )  # new default name should be onnx_graphsurgeon_node_3
        assert graph.nodes[-1].name == "onnx_graphsurgeon_node_3"

    def test_layer_with_attrs(self):
        graph = Graph()
        outputs = graph.layer(op="Add", name="node", attrs={"fake_attr": 0})
        assert len(graph.nodes) == 1
        assert graph.nodes[-1].op == "Add"
        assert graph.nodes[-1].name == "node"
        assert graph.nodes[-1].attrs["fake_attr"] == 0

    def test_layer_with_tensors(self):
        x0 = Variable("x0")
        x1 = Variable("x1")
        y0 = Variable("y0")
        y1 = Variable("y1")
        graph = Graph()

        outputs = graph.layer(op="Fake", inputs=[x0, x1], outputs=[y0, y1])
        assert outputs == [y0, y1]
        assert len(graph.nodes) == 1
        assert graph.nodes[-1].inputs == [x0, x1]
        assert graph.nodes[-1].outputs == outputs

    def test_layer_with_strings(self):
        x0 = "x0"
        x1 = "x1"
        y0 = "y0"
        y1 = "y1"
        graph = Graph()

        outputs = graph.layer(op="Fake", inputs=[x0, x1], outputs=[y0, y1])
        assert len(graph.nodes) == 1
        assert [
            prefix in tensor.name
            for prefix, tensor in zip([x0, x1], graph.nodes[-1].inputs)
        ]
        assert [
            prefix in tensor.name
            for prefix, tensor in zip([y0, y1], graph.nodes[-1].outputs)
        ]
        assert graph.nodes[-1].outputs == outputs

    def test_layer_with_arrays(self):
        x0 = np.array([1])
        x1 = np.array([1])
        y0 = "y0"
        y1 = "y1"
        graph = Graph()

        outputs = graph.layer(op="Fake", inputs=[x0, x1], outputs=[y0, y1])
        assert [
            prefix in tensor.name
            for prefix, tensor in zip([y0, y1], graph.nodes[-1].outputs)
        ]
        assert len(graph.nodes) == 1
        assert graph.nodes[-1].inputs[0].values == x0
        assert graph.nodes[-1].inputs[1].values == x1
        assert graph.nodes[-1].outputs == outputs

    def test_layer_with_iterables(self):
        x0 = [1]
        x1 = (1,)
        y0 = "y0"
        y1 = "y1"
        graph = Graph()

        outputs = graph.layer(op="Fake", inputs=[x0, x1], outputs=[y0, y1])
        assert [
            prefix in tensor.name
            for prefix, tensor in zip([y0, y1], graph.nodes[-1].outputs)
        ]
        assert len(graph.nodes) == 1
        assert graph.nodes[-1].inputs[0].values == x0
        assert graph.nodes[-1].inputs[1].values == x1
        assert graph.nodes[-1].outputs == outputs


class TestFunctionCall(object):
    def make_graph_and_func(self):
        func_output = Variable("test_output", shape=[1, 2, 3], dtype=np.int32)
        func = Function(
            "TestFunction", inputs=[Variable("test_input")], outputs=[func_output]
        )
        graph = Graph(functions=[func])
        return graph, func_output

    def check_outputs_match(self, func_output, node_outputs):
        assert len(node_outputs) == 1
        output = node_outputs[0]

        assert output is not func_output
        assert output.name
        assert output.shape == func_output.shape
        assert output.dtype == func_output.dtype

    def test_function_default_outputs(self):
        graph, func_output = self.make_graph_and_func()

        # No outputs given, but they should be created.
        outputs = graph.TestFunction(inputs=["input"])
        self.check_outputs_match(func_output, outputs)

    def test_function_default_outputs_string_names(self):
        graph, func_output = self.make_graph_and_func()

        # Output name given, it should be preserved.
        outputs = graph.TestFunction(inputs=["input"], outputs=["output"])
        self.check_outputs_match(func_output, outputs)
        assert outputs[0].name.startswith("output")

    def test_function_default_outputs_existing_tensor(self):
        graph, _ = self.make_graph_and_func()

        # Output tensor provided, it should not be changed.
        existing_tensor = Variable("output", shape=[2, 3, 4], dtype=np.float32)
        outputs = graph.TestFunction(inputs=["input"], outputs=[existing_tensor])
        assert outputs[0] is existing_tensor


def tensors_linear_graph():
    inputs = [Variable(name="x")]
    intermediate0 = Variable(name="intermediate0")
    intermediate1 = Variable(name="intermediate1")
    intermediate2 = Variable(name="intermediate2")
    outputs = [Variable(name="y")]

    tensors = inputs + [intermediate0, intermediate1, intermediate2] + outputs
    tensors = {tensor.name: tensor for tensor in tensors}
    # Nodes are NOT in topo order.
    nodes = [
        Node(op="Add", name="Test0", inputs=inputs, outputs=[intermediate0]),
        Node(op="Add", name="Test1", inputs=[intermediate0], outputs=[intermediate1]),
        Node(op="Add", name="Test2", inputs=[intermediate1], outputs=[intermediate2]),
        Node(op="Add", name="Test3", inputs=[intermediate2], outputs=outputs),
    ]
    return Graph(nodes=nodes, inputs=inputs, outputs=outputs), nodes, tensors


class TestTensors(object):
    # Calling `tensors()` should not modify tensors in the graph.
    def test_tensors_does_not_modify_tensors(self):
        graph, _, _ = tensors_linear_graph()
        graph_tensors = graph.tensors()
        # Generate a new graph to compare against
        _, _, tensors = tensors_linear_graph()

        assert set(tensors.keys()) == set(graph_tensors.keys())

        for name, tensor in tensors.items():
            graph_tensor = graph_tensors[name]
            assert tensor == graph_tensor
            assert tensor.inputs == graph_tensor.inputs
            assert tensor.outputs == graph_tensor.outputs

    # Check that tensors includes tensors not attached to nodes
    def test_tensors_includes_non_node_tensors(self):
        X = Constant("X", values=np.ones(shape=(64, 64), dtype=np.float32))
        graph = Graph(inputs=[], outputs=[X])
        tensor_map = graph.tensors()
        assert "X" in tensor_map
        assert tensor_map["X"] == X

    def test_tensors_check_duplicates(self):
        inputs = [Variable(name="x")]
        outputs = [Variable(name="x")]  # Distinct tensors with the same name
        nodes = [
            Node(op="Add", name="Test", inputs=inputs, outputs=outputs),
        ]
        graph = Graph(nodes=nodes, inputs=inputs, outputs=outputs)

        with pytest.raises(OnnxGraphSurgeonException):
            graph.tensors(check_duplicates=True)

    def test_tensors_with_duplicates_check_disabled(self):
        inputs = [Variable(name="x")]
        outputs = [Variable(name="x")]  # Distinct tensors with the same name
        nodes = [
            Node(op="Add", name="Test", inputs=inputs, outputs=outputs),
        ]
        graph = Graph(nodes=nodes, inputs=inputs, outputs=outputs)

        # This should *not* throw
        graph.tensors(check_duplicates=False)


def toposort_linear_graph():
    inputs = [Variable(name="x")]
    intermediate0 = Variable(name="intermediate0")
    intermediate1 = Variable(name="intermediate1")
    intermediate2 = Variable(name="intermediate2")
    outputs = [Variable(name="y")]
    # Nodes are NOT in topo order.
    nodes = [
        Node(op="Add", name="Test0", inputs=inputs, outputs=[intermediate0]),
        Node(op="Add", name="Test2", inputs=[intermediate1], outputs=[intermediate2]),
        Node(op="Add", name="Test3", inputs=[intermediate2], outputs=outputs),
        Node(op="Add", name="Test1", inputs=[intermediate0], outputs=[intermediate1]),
    ]
    expected_node_order = [nodes[0], nodes[3], nodes[1], nodes[2]]
    return Graph(nodes=nodes, inputs=inputs, outputs=outputs), expected_node_order


# Graph structure:
#  x
#  |
#  Test0
#  |
# intermediate0
#   |           \
# Test1         Test2
#   |              |
#  intermediate1   intermediate2
#              \    /
#              Test3
#                |
#                y
def toposort_diamond_graph():
    inputs = [Variable(name="x")]
    intermediate0 = Variable(name="intermediate0")
    intermediate1 = Variable(name="intermediate1")
    intermediate2 = Variable(name="intermediate2")
    outputs = [Variable(name="y")]
    # Nodes are NOT in topo order:
    nodes = [
        Node(op="Add", name="Test0", inputs=inputs, outputs=[intermediate0]),
        # If the end of the diamond shape occurs before the branches, it can trigger a path
        # where a false cycle is detected - this test guards against that.
        Node(
            op="Add",
            name="Test3",
            inputs=[intermediate1, intermediate2],
            outputs=outputs,
        ),
        Node(op="Add", name="Test1", inputs=[intermediate0], outputs=[intermediate1]),
        Node(op="Add", name="Test2", inputs=[intermediate0], outputs=[intermediate2]),
    ]
    expected_node_order = [nodes[0], nodes[2], nodes[3], nodes[1]]
    return Graph(nodes=nodes, inputs=inputs, outputs=outputs), expected_node_order


# Graph structure:
# x
# |
# Test0 -> out0 (graph output)
# |
# out0
# |
# Test1 -> out1 (graph output)
# |
# out1
# |
# Test2 -> out2 (graph_output)
def toposort_multi_tier_output_graph():
    inputs = [Variable(name="x")]
    outputs = [Variable(name="out0"), Variable(name="out1"), Variable(name="out2")]
    out0, out1, out2 = outputs
    nodes = [
        Node(op="Add", name="Test2", inputs=[out1], outputs=[out2]),
        Node(op="Add", name="Test0", inputs=inputs, outputs=[out0]),
        Node(op="Add", name="Test1", inputs=[out0], outputs=[out1]),
    ]
    expected_node_order = [nodes[1], nodes[2], nodes[0]]
    return Graph(nodes=nodes, inputs=inputs, outputs=outputs), expected_node_order


# Graph structure:
# x2  x1
# |   |
# Test0
# |
# int0  x0
# |    /
# Test1
# |
# int1  x3
# |    /
# Test2 -> out (graph_output)
def toposort_multi_tier_input_graph():
    inputs = [
        Variable(name="x0"),
        Variable(name="x1"),
        Variable(name="x2"),
        Variable(name="x3"),
    ]
    int0, int1 = [Variable(name="intermediate0"), Variable(name="intermediate1")]
    outputs = [Variable(name="out")]
    x0, x1, x2, x3 = inputs
    nodes = [
        Node(op="Add", name="Test2", inputs=[int1, x3], outputs=outputs),
        Node(op="Add", name="Test0", inputs=[x2, x1], outputs=[int0]),
        Node(op="Add", name="Test1", inputs=[int0, x0], outputs=[int1]),
    ]
    expected_node_order = [nodes[1], nodes[2], nodes[0]]
    return Graph(nodes=nodes, inputs=inputs, outputs=outputs), expected_node_order


# Graph structure:
# x0
# |
# Add
# |
# x1
# |
# Relu
# |
# x2
#
# If:
#   Then:
#          x1  x2
#           |  |
#           Add
#            |
#           res
#   Else:
#          x1  x2
#           |  |
#           Add
#            |
#           res
#  |
# out
#
# In this graph, the subgraph of If implicitly depends on x1/x2 from the outer graph, so the parent If
# node must come after the outer Add/ReLU nodes.
# If we fail to consider such implicit inputs, the If will remain the first node.
def toposort_implicit_subgraph_inputs_graph():
    def make_var(name):
        return Variable(name, shape=(1, 1), dtype=np.float32)

    # Main graph
    inputs = [make_var("x0")]
    const = Constant(name="const", values=np.array([[1.5]], dtype=np.float32))
    cond = Constant(name="cond", values=np.array([True]))
    x1, x2 = [make_var("x1"), make_var("x2")]
    outputs = [make_var("out")]

    # Subgraphs for If
    subgraph_outputs = [make_var("res")]
    subgraph_nodes = [
        Node(op="Add", name="SubgraphTest0", inputs=[x1, x2], outputs=subgraph_outputs)
    ]
    subgraph = Graph(nodes=subgraph_nodes, outputs=subgraph_outputs)

    nodes = [
        Node(
            op="If",
            name="Test2",
            inputs=[cond],
            outputs=outputs,
            attrs={"then_branch": subgraph, "else_branch": subgraph},
        ),
        Node(op="Relu", name="Test1", inputs=[x1], outputs=[x2]),
        Node(op="Add", name="Test0", inputs=inputs + [const], outputs=[x1]),
    ]
    expected_node_order = [nodes[2], nodes[1], nodes[0]]
    return Graph(nodes=nodes, inputs=inputs, outputs=outputs), expected_node_order


TOPOSORT_TEST_CASES = [
    toposort_linear_graph,
    toposort_diamond_graph,
    toposort_multi_tier_output_graph,
    toposort_multi_tier_input_graph,
    toposort_implicit_subgraph_inputs_graph,
]


class TestToposort(object):
    @staticmethod
    def make_single_node_function(name: str, dependencies: List[str]) -> Function:
        """
        Create a function which uses all the functions given in 'dependencies'.
        """
        func = Function(name, inputs=[Variable("input")], outputs=[Variable("output")])
        if not dependencies:
            return func
        intermediate = func.inputs[0]
        for i, dep in enumerate(dependencies):
            new_intermediate = Variable(f"inter_{i}")
            func.nodes.append(
                Node(
                    dep,
                    domain=Function.DEFAULT_DOMAIN,
                    inputs=[intermediate],
                    outputs=[new_intermediate],
                )
            )
            intermediate = new_intermediate
        func.outputs = [intermediate]
        return func

    @pytest.mark.parametrize("toposort_test_case", TOPOSORT_TEST_CASES)
    def test_topologically_sort(self, toposort_test_case):
        graph, expected_node_order = toposort_test_case()
        assert graph.nodes != expected_node_order
        graph.toposort()
        assert graph.nodes == expected_node_order

    @pytest.mark.parametrize("toposort_test_case", TOPOSORT_TEST_CASES)
    def test_toposort_nested(self, toposort_test_case):
        subgraph, expected_node_order = toposort_test_case()
        assert subgraph.nodes != expected_node_order

        # Wrap the graph within a subgraph
        inp = Variable("input")
        id_out = Variable("id_out")
        identity = Node(op="Identity", inputs=[inp], outputs=[id_out])

        # Make the subgraph take an input from the outer graph node
        # If toposort tries to take the node id, it'll fault.
        subgraph.nodes[0].inputs.append(id_out)

        out = Variable("output")
        nested = Node(
            op="Nested", inputs=[id_out], outputs=[out], attrs={"subgraph": subgraph}
        )

        graph = Graph(nodes=[identity, nested], inputs=[inp], outputs=[out])
        graph.toposort(recurse_subgraphs=True)

        assert subgraph.nodes == expected_node_order

    def test_function(self):
        graph, expected_node_order = toposort_multi_tier_input_graph()
        func = Function(
            "Test", nodes=graph.nodes, inputs=graph.inputs, outputs=graph.outputs
        )
        func.toposort()
        assert func.nodes == expected_node_order

    def test_function_order(self):
        # Check that toposort re-orders functions in topological order.
        func1 = self.make_single_node_function("func1", [])
        func2 = self.make_single_node_function("func2", ["func1"])
        func3 = self.make_single_node_function("func3", ["func2"])
        func4 = self.make_single_node_function("func4", ["func3", "func2"])
        funcs = [func3, func2, func4, func1]
        graph = Graph(functions=funcs)
        graph.toposort()
        assert graph.functions == [func1, func2, func3, func4]

    def test_function_circular_dep_simple(self):
        func = self.make_single_node_function("func", ["func"])
        graph = Graph(functions=[func])
        try:
            graph.toposort()
            assert False, "Should have raised"
        except OnnxGraphSurgeonException:
            pass

    def test_function_circular_dep_complicated(self):
        # Circular dependency [func2 -> func3 -> func4 -> func2]
        func1 = self.make_single_node_function("func1", ["func2", "func5"])
        func2 = self.make_single_node_function("func2", ["func3"])
        func3 = self.make_single_node_function("func3", ["func4", "func5"])
        func4 = self.make_single_node_function("func4", ["func2", "func5"])
        func5 = self.make_single_node_function("func5", [])
        funcs = [func3, func2, func4, func1, func5]
        graph = Graph(functions=funcs)
        try:
            graph.toposort()
            assert False, "Should have raised"
        except OnnxGraphSurgeonException:
            pass


def build_basic_graph():
    inputs = [Variable(name="x")]
    outputs = [Variable(name="y")]
    nodes = [
        Node(op="Add", name="Test", inputs=inputs, outputs=outputs),
    ]
    return Graph(nodes=nodes, inputs=inputs, outputs=outputs, ir_version=10)


def build_two_layer_graph():
    inputs = [Variable(name="x")]
    intermediate_tensor = Variable(name="intermediate")
    outputs = [Variable(name="y")]
    nodes = [
        Node(op="Add", name="Test0", inputs=inputs, outputs=[intermediate_tensor]),
        Node(op="Add", name="Test1", inputs=[intermediate_tensor], outputs=outputs),
    ]
    return Graph(nodes=nodes, inputs=inputs, outputs=outputs, ir_version=10)


def build_two_layer_graph_multiple_io():
    inputs = [Variable(name="x0"), Variable(name="x1")]
    intermediate_tensor = Variable(name="intermediate")
    outputs = [Variable(name="y0"), Variable(name="y1")]
    nodes = [
        Node(op="Add", name="Test0", inputs=inputs, outputs=[intermediate_tensor]),
        Node(op="Add", name="Test1", inputs=[intermediate_tensor], outputs=outputs),
    ]
    return Graph(nodes=nodes, inputs=inputs, outputs=outputs, ir_version=10)


def build_function_with_unused_node():
    func = Function("Test")

    A = Variable("A", dtype=np.float32, shape=(1, 1))
    B = Variable("B", dtype=np.float32, shape=(1, 1))
    X = Variable("X", dtype=np.float32, shape=(1, 1))
    Y = Variable("Y", dtype=np.float32, shape=(1, 1))

    func.inputs = [A, X]
    func.outputs = [B]
    func.nodes = [
        Node(op="Sin", inputs=[X], outputs=[Y]),  # this node is unused
        Node(op="Cos", inputs=[A], outputs=[B]),
    ]
    return func


CLEANUP_TEST_CASES = [
    build_basic_graph(),
    build_two_layer_graph(),
    build_two_layer_graph_multiple_io(),
]


class TestCleanup(object):
    @pytest.mark.parametrize("graph", CLEANUP_TEST_CASES)
    def test_get_used_node_ids(self, graph):
        graph_used_nodes = copy.copy(graph.nodes)
        graph_used_tensors = copy.copy(list(graph.tensors().values()))

        unused_tensor = Variable(name="Unused")
        unused_node = Node(
            op="Unused", inputs=[graph.inputs[0]], outputs=[unused_tensor]
        )
        graph.nodes.append(unused_node)

        with graph.node_ids():
            used_node_ids, used_tensors = graph._get_used_node_ids()
            assert len(used_node_ids) == len(graph.nodes) - 1
            assert all([node.id in used_node_ids for node in graph_used_nodes])
            assert unused_node.id not in used_node_ids
            assert unused_tensor not in used_tensors
            assert all(
                [used_tensor in used_tensors for used_tensor in graph_used_tensors]
            )

    def test_multi_tier(self):
        graph, _ = toposort_multi_tier_output_graph()
        tensor = graph.outputs.pop()
        unused_node = tensor.inputs[0]
        graph.cleanup()  # Should remove just the Test2 node as out1 is still an output.
        assert unused_node not in graph.nodes
        assert len(graph.nodes) == 2
        assert len(graph.outputs) == 2

        tensor_map = graph.tensors()
        assert tensor.name not in tensor_map

    def test_remove_unused_node_outputs(self):
        graph, _ = toposort_linear_graph()
        graph.toposort()
        graph_output = graph.outputs[0]

        dummy = Variable("dummy")
        # Add unused tensor to a node in the middle of the graph.
        # Since it does not contribute to graph outputs, it should be removed.
        graph.nodes[1].outputs.append(dummy)

        graph.cleanup(remove_unused_node_outputs=True)
        assert dummy not in graph.nodes[1].outputs
        assert graph.outputs[0] == graph_output  # Graoh outputs will never be removed

    def test_graph_input_producers(self):
        graph, _ = toposort_linear_graph()
        tensor_map = graph.tensors()
        assert "x" in tensor_map

        graph.inputs = [tensor_map["intermediate0"]]

        graph.cleanup()
        cleaned_tensor_map = graph.tensors()
        assert "x" not in cleaned_tensor_map

    @pytest.mark.parametrize("remove_unused_graph_inputs", [True, False])
    def test_independent_path(self, remove_unused_graph_inputs):
        graph, _ = toposort_linear_graph()
        # Build out a path totally unrelated to rest of the graph
        indep0 = Variable(name="indep0")
        indep1 = Variable(name="indep1")
        node = Node(op="IndepTest", inputs=[indep0], outputs=[indep1])
        graph.nodes.append(node)
        graph.inputs.append(indep0)
        graph.cleanup(remove_unused_graph_inputs=remove_unused_graph_inputs)
        assert indep0 not in graph.inputs or not remove_unused_graph_inputs
        assert node not in graph.nodes or not remove_unused_graph_inputs

        tensor_map = graph.tensors()
        assert indep0.name not in tensor_map or not remove_unused_graph_inputs
        assert indep1.name not in tensor_map or not remove_unused_graph_inputs

    def test_nested_graph(self, nested_graph):
        nested_node = nested_graph.nodes[1]
        nested_inp = nested_node.inputs[0]
        nested_out = nested_node.outputs[0]
        subgraph = nested_node.attrs["body"]

        assert "id_out" in nested_graph.tensors()
        nested_graph.cleanup(recurse_subgraphs=True)
        # Clean up should not remove a tensor whose only output node is a subgraph.
        assert "id_out" in nested_graph.tensors()

        # Clean up should not modify the nested nodes inputs or outputs
        assert nested_node.inputs == [nested_inp]
        assert nested_node.outputs == [nested_out]

        # Next we'll clean up the subgraph by recursing from the top-level
        assert subgraph.nodes
        subgraph.outputs.clear()
        nested_graph.cleanup(recurse_subgraphs=True)
        assert not subgraph.nodes

    def test_node_used_only_in_nested_graph(self):
        X = Variable("X", dtype=np.float32, shape=(1,))
        Y = Variable("Y", dtype=np.float32, shape=(1,))
        graph = Graph(inputs=[X, Y])

        X_p = graph.identity(
            X
        )  # X_p is only used by the subgraph, not in the outer graph.

        subgraph_inp = Variable("subgraph_input", dtype=np.float32, shape=(1,))
        subgraph = Graph(inputs=[subgraph_inp])
        subgraph.outputs = [subgraph.add(subgraph_inp, X_p)]

        graph.outputs = [graph.nested(Y, subgraph)]

        graph.cleanup(remove_unused_graph_inputs=True)

        assert graph.nodes[0].op == "Identity"
        assert graph.nodes[0].inputs == [X]

    def test_input_is_output(self):
        graph = Graph()

        A = Variable("A", dtype=np.float32, shape=(1, 1))
        B = Variable("B", dtype=np.float32, shape=(1, 1))

        C = graph.add(A, B)

        graph.inputs = [A, B]
        graph.outputs = [C, B, A]  # Out of order w/ respect to Add node inputs

        # Graph should remain unchanged after cleanup, including I/O tensors.
        graph.cleanup()

        assert graph.inputs == [A, B]
        assert graph.outputs == [C, B, A]
        assert len(graph.nodes) == 1
        assert graph.nodes[0].inputs == [A, B]
        assert graph.nodes[0].outputs == [C]

    def test_function(self):
        func = build_function_with_unused_node()

        func.cleanup()
        assert {i.name for i in func.inputs} == {"A", "X"}
        assert {o.name for o in func.outputs} == {"B"}
        assert len(func.nodes) == 1

    def test_graph_cleans_up_function(self):
        graph = Graph()
        func = build_function_with_unused_node()
        graph.functions.append(func)

        # Cleaning up the graph should by default also cleanup the function.
        graph.cleanup()
        assert {i.name for i in func.inputs} == {"A", "X"}
        assert {o.name for o in func.outputs} == {"B"}
        assert len(func.nodes) == 1


class TestCopy(object):
    def test_basic(self):
        graph = Graph(
            nodes=[Node(op="Test")],
            inputs=[Variable("test")],
            outputs=[Variable("test")],
            name="test-name",
            doc_string="test-docstring",
            import_domains=["fake-import-domain"],
            opset=-1,
        )
        new_graph = graph.copy()

        assert new_graph == graph
        assert new_graph.nodes == graph.nodes
        assert new_graph.inputs == graph.inputs
        assert new_graph.outputs == graph.outputs
        assert new_graph.name == graph.name
        assert new_graph.doc_string == graph.doc_string
        assert new_graph.import_domains == graph.import_domains
        assert new_graph.opset == graph.opset

    def test_copy(self):
        def make_graph():
            graph, _ = toposort_multi_tier_output_graph()
            graph.outputs.pop()
            # Deep copy should work with empty tensors
            graph.nodes[0].inputs.append(Variable.empty())
            graph.nodes[0].outputs.append(Variable.empty())
            return graph

        graph = make_graph()
        new_graph = graph.copy()
        assert graph == new_graph

        # Running cleanup on the first graph should not affect the copy
        graph.cleanup()
        assert graph != new_graph
        assert new_graph == make_graph()

    def test_copy_with_subgraph(self, nested_graph):
        new_graph = nested_graph.copy()
        assert new_graph == nested_graph

        new_subgraph = new_graph.nodes[1].attrs["body"]

        id_out = new_subgraph.nodes[0].inputs[0]
        assert id_out.name == "id_out"
        assert len(id_out.inputs) == 1
        assert id_out.inputs[0].op == "Identity"
        assert id_out.inputs[0].inputs[0].name == "input"

        new_subgraph.nodes[0].outputs.clear()
        new_subgraph.nodes[1].inputs.clear()

        subgraph = nested_graph.nodes[1].attrs["body"]
        assert subgraph.nodes[0].outputs
        assert subgraph.nodes[1].inputs

        new_graph.outputs.clear()
        new_graph.cleanup()

        assert nested_graph.outputs
        assert len(nested_graph.nodes) == 2
        assert len(subgraph.nodes) == 2

    # If the subgraph has a tensor with the same name as the outer graph,
    # the subgraph copy should include a copy of the subgraph tensor, not the outer
    # graph tensor.
    def test_copy_with_subgraph_dup_tensors(self):
        inp = Variable("input", dtype=np.float32, shape=(4, 5))
        graph = Graph(inputs=[inp])

        # We'll use shape to distinguish inner/outer tensor
        subgraph_inp = Variable("input", dtype=np.float32, shape=(1, 2))
        subgraph = Graph(inputs=[subgraph_inp])

        graph.outputs = [graph.nested(inp, subgraph)]

        graph_copy = graph.copy()
        assert graph_copy.nodes[0].attrs["body"].inputs[0].shape == (1, 2)

    def test_copy_with_subgraph_dup_const_tensors(self):
        inp = Constant("input", values=np.ones(dtype=np.float32, shape=(4, 5)))
        graph = Graph()

        # We'll use shape to distinguish inner/outer tensor
        subgraph_inp = Constant("input", values=np.ones(dtype=np.float32, shape=(1, 2)))
        subgraph = Graph()
        subgraph.outputs = [subgraph.identity(subgraph_inp)]

        graph.outputs = [graph.nested(inp, subgraph)]

        graph_copy = graph.copy()
        assert graph_copy.nodes[0].attrs["body"].nodes[0].inputs[0].shape == (1, 2)

    def test_function(self):
        func = Function(
            "Test",
            domain="onnx-gs.test",
            nodes=[Node(op="Add")],
            inputs=[Variable("input")],
            outputs=[Variable("output")],
            doc_string="docstring",
            opset=15,
            import_domains=["test"],
            attrs={"attr1": None, "attr2": np.array([1, 2, 3])},
        )
        func_copy = func.copy()
        assert func.name == func_copy.name
        assert func.domain == func_copy.domain
        assert func.nodes == func_copy.nodes
        assert func.inputs == func_copy.inputs
        assert func.outputs == func_copy.outputs
        assert func.doc_string == func_copy.doc_string
        assert func.opset == func_copy.opset
        assert func.import_domains == func_copy.import_domains
        assert func.attrs["attr1"] == func_copy.attrs["attr1"]
        assert np.all(func.attrs["attr2"] == func_copy.attrs["attr2"])

        assert func.nodes is not func_copy.nodes
        assert func.inputs is not func_copy.inputs
        assert func.outputs is not func_copy.outputs
        assert func.attrs is not func_copy.attrs

        assert func.nodes[0] is not func_copy.nodes[0]
        assert func.inputs[0] is not func_copy.inputs[0]
        assert func.outputs[0] is not func_copy.outputs[0]
        assert func.attrs["attr2"] is not func_copy.attrs["attr2"]


@pytest.fixture
def simple_foldable():
    # Graph:
    # c = (a + b)
    # output = input + c
    # Should fold to:
    # output = input + c
    weights = np.ones(shape=(1, 3), dtype=np.float32)

    graph = Graph(ir_version=10)
    inp = Variable("input", shape=(1, 3), dtype=np.float32)
    c = graph.add(weights, weights, name="c")
    out = graph.add(inp, c, name="out")

    graph.inputs = [inp]
    graph.outputs = [out]
    yield graph


@pytest.fixture
def one_hop_foldable():
    # Graph:
    # c = (a + b)
    # e = (c + d)
    # output = input + e
    # Should fold to:
    # output = input + e
    weights = np.ones(shape=(1, 3), dtype=np.float32)

    graph = Graph(ir_version=10)
    inp = Variable("input", shape=(1, 3), dtype=np.float32)
    c = graph.add(weights, weights, name="c")
    e = graph.add(c, weights, name="e")
    out = graph.add(inp, e)

    graph.inputs = [inp]
    graph.outputs = [out]
    yield graph


@pytest.fixture
def foldable_with_invalid_node():
    # Graph
    # c = (a + b)
    # e = fake(d)
    # f = (e + c)
    # out = inp + f
    #
    # c should be folded even though e is the output of an
    # invalid node.
    weights = np.ones(shape=(1, 3), dtype=np.float32)

    graph = Graph(ir_version=10)
    inp = Variable("input", shape=(1, 3), dtype=np.float32)
    c = graph.add(weights, weights, name="c")
    e = graph.fake(weights, name="e")
    f = graph.add(e, c, name="f")
    out = graph.add(inp, f, name="output")

    graph.inputs = [inp]
    graph.outputs = [out]
    yield graph


@pytest.fixture
def foldable_with_local_functions():
    dtype = np.float32
    counter = 0

    def const():
        nonlocal counter
        counter += 1
        return Constant(f"constant_{counter}", np.ones(1, dtype=np.float32))

    func_inner = Function("FuncInner", ir_version=10)
    func_outer = Function("FuncOuter", ir_version=10)
    funcs = [func_inner, func_outer]

    # func_inner(x) = x + 1
    func_inner.inputs = [Variable("input", dtype=dtype)]
    func_inner.outputs = [Variable("output", dtype=dtype)]
    func_inner.nodes = [
        Node(
            "Add",
            inputs=[func_inner.inputs[0], const()],
            outputs=[func_inner.outputs[0]],
        )
    ]

    # func_outer(x) = func_inner(1) * x
    func_outer.inputs = [Variable("input", dtype=dtype)]
    func_outer.functions = [func_inner]
    func_outer_intermediate = func_outer.FuncInner(inputs=[const()])[0]
    func_outer.outputs = [func_outer.mul(func_outer.inputs[0], func_outer_intermediate)]

    # a = func_inner(input)
    # b = func_outer(input)
    # c = 1 + 1
    # d = a + b
    # e = c + d
    # output = e + c
    graph = Graph(
        inputs=[Variable("graph_input", dtype=dtype)], functions=funcs, ir_version=10
    )
    var0 = graph.FuncInner(inputs=[const()])[0]
    var1 = graph.FuncOuter(inputs=[const()])[0]
    var2 = graph.add(const(), const())
    var3 = graph.add(var0, var1)
    var4 = graph.add(var2, var3)
    graph.outputs = [graph.add(var2, var4)]
    graph.outputs[0].dtype = dtype

    yield graph


class TestFoldConstants(object):
    @pytest.mark.parametrize("partitioning", [None, "basic", "recursive"])
    def test_basic(self, simple_foldable, partitioning):
        inp = simple_foldable.inputs[0]

        simple_foldable.fold_constants(partitioning=partitioning).cleanup(
            remove_unused_graph_inputs=True
        )

        # Extra node should be removed
        assert len(simple_foldable.nodes) == 1
        assert simple_foldable.nodes[0].inputs[0] == inp
        assert simple_foldable.nodes[0].inputs[1].name == "c"

        # Value should be computed correctly
        assert np.all(
            simple_foldable.nodes[0].inputs[1].values
            == np.ones(shape=(1, 3), dtype=np.float32) * 2
        )

    def test_one_hop(self, one_hop_foldable):
        inp = one_hop_foldable.inputs[0]

        one_hop_foldable.fold_constants().cleanup()

        # Extra nodes should be removed
        assert len(one_hop_foldable.nodes) == 1
        assert one_hop_foldable.nodes[0].inputs[0] == inp
        assert one_hop_foldable.nodes[0].inputs[1].name == "e"

        # Value should be computed correctly
        assert np.all(
            one_hop_foldable.nodes[0].inputs[1].values
            == np.ones(shape=(1, 3), dtype=np.float32) * 3
        )

    def test_with_invalid_nodes(self, foldable_with_invalid_node):
        foldable_with_invalid_node.fold_constants(partitioning="recursive").cleanup()

        tensor_map = foldable_with_invalid_node.tensors()

        assert len(foldable_with_invalid_node.nodes) == 3
        assert foldable_with_invalid_node.nodes[0].op == "Fake"
        assert foldable_with_invalid_node.nodes[1].op == "Add"
        assert foldable_with_invalid_node.nodes[2].op == "Add"
        assert np.all(
            tensor_map["c"].values == (np.ones(shape=(1, 3), dtype=np.float32) * 2)
        )

    def test_with_invalid_nodes_no_recursive(self, foldable_with_invalid_node):
        # No folding should take place without recursive partitioning
        original = foldable_with_invalid_node.copy()
        assert foldable_with_invalid_node.fold_constants() == original

    def test_no_foldable_constants(self):
        inp0 = Variable("input0", shape=(1, 3), dtype=np.float32)
        inp1 = Variable("input1", shape=(1, 3), dtype=np.float32)
        out = Variable("output", shape=(1, 3), dtype=np.float32)

        nodes = [Node("Add", inputs=[inp0, inp1], outputs=[out])]

        graph = Graph(nodes=nodes, inputs=[inp0, inp1], outputs=[out], ir_version=10)

        graph.fold_constants().cleanup()

        assert len(graph.nodes) == 1
        assert graph.nodes[0].inputs == [inp0, inp1]

    def test_const_node(self):
        graph = Graph(ir_version=10)
        values = np.ones((1, 3, 3), dtype=np.int64)
        graph.outputs = [graph.constant(values=values)]

        assert isinstance(graph.outputs[0], Variable)

        graph.fold_constants().cleanup()

        assert isinstance(graph.outputs[0], Constant)
        assert np.all(graph.outputs[0].values == values)
        assert not graph.nodes

    def test_shape_of_constant_tensor(self):
        graph = Graph(ir_version=10)
        values = np.ones((1, 3, 3), dtype=np.int64)
        const = Constant("const", values=values)
        graph.outputs = [graph.shape(const)]

        graph.fold_constants().cleanup()

        assert not graph.nodes
        assert isinstance(graph.outputs[0], Constant)
        assert np.all(graph.outputs[0].values == (1, 3, 3))

    def test_shape_of_constant_node(self):
        graph = Graph(ir_version=10)
        values = np.ones((1, 3, 3), dtype=np.int64)
        const = graph.constant(values=values)
        graph.outputs = [graph.shape(const)]

        graph.fold_constants().cleanup()

        assert not graph.nodes
        assert isinstance(graph.outputs[0], Constant)
        assert np.all(graph.outputs[0].values == (1, 3, 3))

    # Cannot fold shape nodes if they have dynamically shaped inputs.
    def test_shape_of_variable_tensor_dynamic_shape(self):
        var = Variable("var", dtype=np.float32, shape=("", -1, 0, 4))
        graph = Graph(inputs=[var], ir_version=10)
        graph.outputs = [graph.shape(var)]

        graph.fold_constants().cleanup()

        assert len(graph.nodes) == 1
        assert graph.nodes[0].op == "Shape"
        assert isinstance(graph.outputs[0], Variable)

    def test_shape_of_variable_tensor_static_shape(self):
        var = Variable("var", dtype=np.float32, shape=(1, 3, 4))
        graph = Graph(inputs=[var], ir_version=10)
        graph.inputs = [var]
        graph.outputs = [graph.shape(var)]

        graph.fold_constants().cleanup()

        assert not graph.nodes
        assert isinstance(graph.outputs[0], Constant)
        assert np.all(graph.outputs[0].values == (1, 3, 4))

    def test_shape_of_variable_tensor_multiple_shapes(self):
        graph = Graph(ir_version=10)
        var = Variable("var", dtype=np.float32, shape=(1, 3, 4))
        var2 = Variable("var2", dtype=np.float32, shape=tuple())  # Scalar
        graph.inputs = [var, var2]
        graph.outputs = [graph.shape(var), graph.identity(var), graph.shape(var2)]

        graph.fold_constants().cleanup()

        assert len(graph.nodes) == 1
        assert graph.nodes[0].op == "Identity"
        assert isinstance(graph.outputs[0], Constant)
        assert np.all(graph.outputs[0].values == (1, 3, 4))
        assert isinstance(graph.outputs[2], Constant)
        assert np.all(graph.outputs[2].values == tuple())

    def test_shape_of_variable_tensor_static_shape_no_fold(self):
        graph = Graph(ir_version=10)
        var = Variable("var", dtype=np.float32, shape=(1, 3, 4))
        graph.inputs = [var]
        graph.outputs = [graph.shape(var)]

        graph.fold_constants(fold_shapes=False).cleanup()

        assert len(graph.nodes) == 1
        assert graph.nodes[0].op == "Shape"
        assert isinstance(graph.outputs[0], Variable)

    # Constant folding should not cause constant tensors in the model to be loaded.
    def test_no_load_constants(self):
        graph = gs.import_onnx(const_foldable().load())

        new_graph = graph.fold_constants()

        def check_no_const_loaded(graph):
            num_lazy_constants = 0
            for tensor in graph.tensors().values():
                if isinstance(tensor, Constant) and isinstance(
                    tensor._values, LazyValues
                ):
                    num_lazy_constants += 1
            assert (
                num_lazy_constants == 3
            )  # Graph starts with 3 constants - none should be loaded.

        check_no_const_loaded(graph)
        check_no_const_loaded(new_graph)

    @pytest.mark.parametrize(
        "shape, indices",
        [
            (("batch", 3, "height", "width"), 1),  # Scalar indices case
            (None, 1),  # Shape not inferered case
            (("batch", 3, "height", "width"), [1]),
            (("batch", 3, "height", 224), [1, 3]),
            (("batch", 3, 224, 224), [1, 2, 3]),
        ],
    )
    def test_shape_gather(self, shape, indices):
        indices = np.array(indices)

        inp = Variable("input", dtype=np.float32, shape=shape)
        graph = Graph(inputs=[inp], ir_version=10)

        inp_shape = graph.shape(inp)
        shape_part = graph.gather(inp_shape, indices=indices)
        graph.outputs = [
            graph.add(shape_part, shape_part),
            graph.gather(inp_shape, indices=[0]),
            graph.gather(inp_shape, indices=np.array(0)),
        ]

        graph.fold_constants()

        if shape is not None:
            assert isinstance(graph.outputs[0], Constant)
            expected_shape = np.array(shape)[indices].astype(np.int64) * 2
            assert np.all(graph.outputs[0].values == expected_shape)
        else:
            assert isinstance(graph.outputs[0], Variable)

        assert isinstance(graph.outputs[1], Variable)
        assert isinstance(graph.outputs[2], Variable)

    @pytest.mark.parametrize(
        "shape, starts, ends, axes, steps, expected",
        [
            (
                ("batch", 3, "height", "width"),
                1,
                2,
                0,
                1,
                [3],
            ),  # Scalar starts/ends case
            (("batch", 3, "height", "width"), [1], [2], [0], [1], [3]),
            (
                ("batch", 3, 5, "width"),
                [1],
                [-1],
                [0],
                [1],
                [3, 5],
            ),  # Negative ends case
            (
                ("batch", 3, 5, 7),
                [1],
                [2000],
                [0],
                [1],
                [3, 5, 7],
            ),  # Past end, ends case
            (("batch", 3, 5, 7), [-2], [4], [0], [1], [5, 7]),  # Negative starts case
            (("batch", 3, 5, 7), [-2], [4], [1], [1], None),  # Non-zero axes case
            (("batch", 3, 5, "width"), [-2], [4], [1], [1], None),  # Dynamic case
            (("batch", 3, 5, 7), [1], [4], [0], [2], [3, 7]),  # Non-one steps case
            (("batch", 3, 5, 7), [4], [0], [0], [-1], [7, 5, 3]),  # Negative steps case
        ],
    )
    def test_shape_slice(self, shape, starts, ends, axes, steps, expected):
        inp = Variable("input", dtype=np.float32, shape=shape)
        graph = Graph(inputs=[inp], ir_version=10)

        inp_shape = graph.shape(inp)
        graph.outputs = [
            graph.slice(
                inp_shape,
                np.array(starts),
                np.array(ends),
                axes=np.array(axes),
                steps=np.array(steps),
            )
        ]

        graph.fold_constants()

        if expected:
            assert isinstance(graph.outputs[0], Constant)
            assert np.all(graph.outputs[0].values == expected)
        else:
            assert isinstance(graph.outputs[0], Variable)

    # In the single input case, we should derive starts/ends/axes/steps from the attributes.
    def test_shape_slice_single_input(self):
        inp = Variable("input", dtype=np.int64, shape=(5, 6, 3, 2))
        graph = Graph(inputs=[inp], ir_version=10)

        inp_shape = graph.shape(inp)
        graph.outputs = [graph.slice(inp_shape)]

        slice_node = graph.outputs[0].inputs[0]

        slice_node.attrs = {
            "axes": [0],
            "starts": [1],
            "ends": [3],
            "steps": [2],
        }

        graph.fold_constants()

        assert isinstance(graph.outputs[0], Constant)
        assert np.all(graph.outputs[0].values == inp.shape[1:3:2])

    def test_with_variable_conditional(self):
        cond = gs.Variable("cond", dtype=bool, shape=(1,))

        X = gs.Variable("X", dtype=np.float32, shape=(1,))
        Y = gs.Constant("Y", values=np.ones((1,), dtype=np.float32))
        graph = Graph(inputs=[X, cond], ir_version=10)

        then_graph = Graph(name="Then", ir_version=10)
        then_graph.outputs = [then_graph.add(Y, Y)]

        else_graph = Graph(name="Else", ir_version=10)
        else_graph.outputs = [else_graph.add(X, else_graph.add(Y, Y))]

        graph.outputs = [graph.if_op(cond, then_graph, else_graph)]

        graph.fold_constants()
        graph.cleanup()

        assert len(then_graph.nodes) == 0
        assert np.all(then_graph.outputs[0].values == (Y.values * 2))

        assert len(else_graph.nodes) == 1
        assert isinstance(else_graph.nodes[0].inputs[1], Constant)
        assert np.all(else_graph.nodes[0].inputs[1].values == (Y.values * 2))

    @pytest.mark.parametrize("cond_value", [True, False])
    @pytest.mark.parametrize("flatten", [True, False])
    def test_flatten_static_conditional(self, flatten, cond_value):
        cond = gs.Constant("cond", values=np.array([cond_value], dtype=bool))

        X = gs.Variable("X", dtype=np.float32, shape=(1,))
        Y = gs.Variable("Y", dtype=np.float32, shape=(1,))
        graph = Graph(inputs=[X, cond], ir_version=10)

        then_graph = Graph(name="Then", ir_version=10)
        then_graph.outputs = [then_graph.relu(then_graph.add(Y, Y))]

        else_graph = Graph(name="Else", ir_version=10)
        else_graph.outputs = [else_graph.add(X, else_graph.add(Y, Y))]

        if_out = graph.if_op(cond, then_graph, else_graph)
        graph.outputs = [if_out]

        graph.fold_constants(flatten_subgraphs=flatten)
        graph.cleanup()

        if flatten:
            assert len(graph.nodes) == 2
            assert graph.nodes[0].op == "Add"
            assert graph.nodes[1].op == "Relu" if cond_value else "Add"

            subgraph = then_graph if cond_value else else_graph
            # Make sure subgraph intermediate tensors are renamed
            assert graph.nodes[0].outputs[0].name == "add_out_0_subg_0_{:}".format(
                subgraph.name
            )
            assert graph.outputs[0].inputs[0] == subgraph.nodes[-1]
            assert subgraph.nodes[-1] == graph.nodes[-1]
        else:
            assert len(graph.nodes) == 1
            assert len(graph.nodes) == 1
            assert graph.nodes[0].op == "If"
            assert graph.outputs[0].inputs[0] == graph.nodes[-1]
        assert graph.outputs == [if_out]

    def test_const_inp_but_non_foldable_nested_graph(self):
        cond = gs.Constant("cond", values=np.array(True))
        X = gs.Variable("X", dtype=np.float32, shape=(1,))

        graph = Graph(inputs=[X], ir_version=10)

        then_graph = Graph(name="Then", ir_version=10)
        then_graph.outputs = [then_graph.add(X, X)]

        else_graph = Graph(name="Else", ir_version=10)
        else_graph.outputs = [else_graph.add(X, else_graph.add(X, X))]

        # Even though if_op looks foldable because it has all constant inputs,
        # it's not, since its subgraphs depend on variables in the outer scope.
        graph.outputs = [graph.if_op(cond, then_graph, else_graph)]

        # This should not raise because the `If` node should be excluded from
        # constant folding.
        graph.fold_constants(error_ok=False, flatten_subgraphs=False).cleanup()

        assert graph.nodes[0].op == "If"
        assert len(then_graph.nodes) == 1
        assert len(else_graph.nodes) == 2

    def test_cast_elision(self):
        graph = gs.import_onnx(shape_cast_elision().load())
        graph.fold_constants().cleanup()
        assert not any(node.op == "Cast" for node in graph.nodes)

    def test_cast_elision_int64(self):
        X = gs.Variable("X", dtype=np.int64, shape=(1,))
        graph = Graph(inputs=[X], ir_version=10)
        casted_x = graph.cast(X, to=onnx.TensorProto.DataType.FLOAT)
        add_out = graph.add(casted_x, casted_x)
        graph.outputs = [graph.cast(add_out, to=onnx.TensorProto.DataType.INT64)]

        graph.fold_constants().cleanup()
        assert graph.nodes[0].op == "Add"

    # Make sure we're lowering constant nodes before running cast elision
    def test_cast_elision_with_constant_node(self):
        inp = gs.Variable("inp", dtype=np.int64, shape=(1,))
        graph = Graph(inputs=[inp], ir_version=10)

        casted_inp = graph.cast(inp, to=onnx.TensorProto.DataType.FLOAT)
        add_out = graph.add(casted_inp, graph.constant(np.array([2], dtype=np.float32)))

        casted_out = graph.cast(add_out, to=onnx.TensorProto.DataType.INT64)
        casted_out.dtype = np.int64

        graph.outputs = [casted_out]

        graph.fold_constants().cleanup()
        assert [node.op for node in graph.nodes] == ["Add"]

        add_const_inp = graph.nodes[0].inputs[1]
        assert isinstance(add_const_inp, Constant)
        assert (
            add_const_inp.dtype == np.int64
        )  # Should have been casted to match dtype of other inputs.

    # For a graph like:
    #
    #     inp
    #      |
    #    Cast
    #      |
    #    Add
    #     |
    #    Cast
    #     |
    #    out
    #
    # 1. We cannot remove the initial `Cast` if it is used outside the `Add` node
    # 2. We cannot perform cast elision at all if the original output of the `Add` node is
    #    used outside the subsequent `Cast` node.
    #
    @pytest.mark.parametrize("use_as_graph_output", [True, False], ids=["graph", ""])
    @pytest.mark.parametrize("use_in_other_node", [True, False], ids=["node", ""])
    # Whether to apply the effects of the first two parameters to the input `Cast` node or to the `Add` node.
    @pytest.mark.parametrize(
        "apply_to_input_cast", [True, False], ids=["input", "output"]
    )
    def test_cast_elision_multi_use_cast(
        self, use_as_graph_output, use_in_other_node, apply_to_input_cast
    ):
        X = gs.Variable("X", dtype=np.int32, shape=(1,))
        graph = Graph(inputs=[X], ir_version=10)
        casted_x = graph.cast(X, to=onnx.TensorProto.DataType.FLOAT)
        add_out = graph.add(casted_x, casted_x)
        uncasted_x = graph.cast(add_out, to=onnx.TensorProto.DataType.INT32)

        graph.outputs = [uncasted_x]

        mutli_use_tensor = casted_x if apply_to_input_cast else add_out
        if use_in_other_node:
            graph.outputs.append(graph.identity(mutli_use_tensor))

        if use_as_graph_output:
            graph.outputs.append(mutli_use_tensor)

        print(graph)
        graph.fold_constants().cleanup()
        ops = [node.op for node in graph.nodes]
        if use_as_graph_output or use_in_other_node:
            if apply_to_input_cast:
                assert graph.nodes[1].inputs[0] == X
                assert graph.nodes[1].outputs[0] == uncasted_x
                assert ops == ["Cast", "Add"] + (
                    ["Identity"] if use_in_other_node else []
                )
            else:
                assert ops == ["Cast", "Add", "Cast"] + (
                    ["Identity"] if use_in_other_node else []
                )
        else:
            assert ops == ["Add"]

    @pytest.mark.parametrize(
        # If layer1_num_bytes is larger than layer0_num_bytes, then it must be a multiple.
        "size_threshold, layer0_num_bytes, layer0_should_fold, layer1_num_bytes, layer1_should_fold",
        [
            # No size threshold - everything should fold.
            (
                None,
                2,
                True,
                4,
                True,
            ),
            # Monotonically increasing but under size threshold - everything should fold.
            (
                8,
                2,
                True,
                4,
                True,
            ),
            # Increasing then decreasing, but under size threshold - everything should fold.
            (
                8,
                2,
                True,
                1,
                True,
            ),
            # All tensors over size threshold - nothing should fold.
            (
                1,
                2,
                False,
                4,
                False,
            ),
            # Second tensor over size threshold - only first tensor should fold.
            (
                3,
                2,
                True,
                4,
                False,
            ),
            # First tensor over size threshold - second tensor should still fold.
            (
                3,
                4,
                False,
                2,
                True,
            ),
        ],
    )
    @pytest.mark.parametrize("push_into_subgraph", [True, False], ids=["subgraph", ""])
    def test_folding_size_threshold(
        self,
        size_threshold,
        layer0_num_bytes,
        layer0_should_fold,
        layer1_num_bytes,
        layer1_should_fold,
        push_into_subgraph,
    ):
        graph = Graph(ir_version=10)

        shape = (1,)

        layer0_repeats = layer0_num_bytes // misc.volume(shape)
        layer0 = graph.tile(np.ones(shape, dtype=np.int8), repeats=[layer0_repeats])
        layer0.inputs[0].name = "Layer0"

        if layer1_num_bytes > layer0_num_bytes:
            layer1_repeats = layer1_num_bytes // layer0_num_bytes
            layer1 = graph.tile(layer0, repeats=[layer1_repeats])
        else:
            layer1 = graph.slice(layer0, starts=[0], ends=[layer1_num_bytes])
        layer1.inputs[0].name = "Layer1"

        graph.outputs = [layer1]

        # Make sure size_threshold option is propagated into subgraphs.
        if push_into_subgraph:
            cond = gs.Variable("cond", dtype=bool, shape=tuple())
            outer_graph = Graph(inputs=[cond])
            outer_graph.if_op(cond, then_graph=graph, else_graph=graph)

            outer_graph.fold_constants(size_threshold=size_threshold)
        else:
            graph.fold_constants(size_threshold=size_threshold)

        # When a tensor is folded, it is disconnected from its producer nodes
        assert len(graph.nodes[0].outputs) == (0 if layer0_should_fold else 1)
        assert len(graph.nodes[1].outputs) == (0 if layer1_should_fold else 1)

    @pytest.mark.parametrize("op", ["Q", "DQ"])
    @pytest.mark.parametrize("add_intermediate_layer", [True, False])
    def test_no_fold_qdq(self, op, add_intermediate_layer):
        dtype = np.float32 if op == "Q" else np.int8
        inp = gs.Constant("input", np.ones(shape=(1, 3, 5, 5), dtype=dtype))
        graph = Graph(inputs=[inp], opset=13, ir_version=10)

        if add_intermediate_layer:
            inp = graph.identity(inp)

        qdq_func = graph.quantize_linear if op == "Q" else graph.dequantize_linear
        graph.outputs = [
            qdq_func(inp, 1.2, np.array(0, dtype=np.int8))
        ]  # Arbitrary scale and zero-point

        graph.fold_constants().cleanup()
        assert len(graph.nodes) == 1
        assert (
            graph.nodes[0].op == "QuantizeLinear" if op == "Q" else "DequantizeLinear"
        )

    @pytest.mark.parametrize(
        "should_exclude_node_func,expected_node_names",
        [
            (
                lambda node: True,
                [
                    "onnx_graphsurgeon_node_1",
                    "onnx_graphsurgeon_node_3",
                    "onnx_graphsurgeon_node_5",
                    "onnx_graphsurgeon_node_7",
                ],
            ),
            (
                lambda node: node.name == "onnx_graphsurgeon_node_5",
                ["onnx_graphsurgeon_node_5", "onnx_graphsurgeon_node_7"],
            ),
            (
                lambda node: node.op == "Add",
                [
                    "onnx_graphsurgeon_node_1",
                    "onnx_graphsurgeon_node_3",
                    "onnx_graphsurgeon_node_5",
                    "onnx_graphsurgeon_node_7",
                ],
            ),
            (
                lambda node: node.op == "Relu",
                [
                    "onnx_graphsurgeon_node_3",
                    "onnx_graphsurgeon_node_5",
                    "onnx_graphsurgeon_node_7",
                ],
            ),
        ],
    )
    def test_custom_should_exclude_node(
        self, should_exclude_node_func, expected_node_names
    ):
        inp = gs.Constant("input", np.ones(shape=(1, 3, 5, 5), dtype=np.float32))
        graph = Graph(inputs=[inp], ir_version=10)

        add_0 = graph.add(inp, inp)  # onnx_graphsurgeon_node_1 -> add_out_0
        relu_0 = graph.relu(add_0)  # onnx_graphsurgeon_node_3 -> relu_out_2
        add_1 = graph.add(relu_0, relu_0)  # onnx_graphsurgeon_node_5 -> add_out_4
        relu_1 = graph.relu(add_1)  # onnx_graphsurgeon_node_7 -> relu_out_6

        graph.outputs = [relu_1]

        graph.fold_constants(should_exclude_node=should_exclude_node_func).cleanup()
        assert [node.name for node in graph.nodes] == expected_node_names

    def test_omitted_optional_inputs_ignored(self):
        # An omitted optional input will show up as a `Variable` with no name.
        # This should *not* prevent us from folding nodes where all other inputs are constants.
        data = gs.Constant("data", np.ones(shape=(3, 5, 5), dtype=np.float32))
        pads = gs.Constant("pads", np.zeros(shape=(6,), dtype=np.int64))
        graph = Graph(ir_version=10)

        pad_0 = graph.pad(data, pads, constant_value=None)
        graph.outputs = [pad_0]

        assert pad_0.inputs[0].inputs[2] == Variable.empty()
        assert len(graph.nodes) == 1

        graph.fold_constants().cleanup()

        assert len(graph.nodes) == 0
        assert isinstance(graph.outputs[0], gs.Constant)

    def test_function(self, simple_foldable):
        func = Function(
            "Test",
            nodes=simple_foldable.nodes,
            inputs=simple_foldable.inputs,
            outputs=simple_foldable.outputs,
            ir_version=10,
        )
        assert len(func.nodes) == 2
        func.fold_constants().cleanup()
        assert len(func.nodes) == 1

    def test_function_with_attributes(self):
        # Nodes that reference function attributes shouldn't be folded.
        input = Variable("input", dtype=np.float32)
        func = Function(
            "Test", inputs=[input], attrs={"softmax_axis": -1}, ir_version=10
        )
        x = func.softmax(
            np.float32([[1, 2, 3]]), axis=Node.AttributeRef("softmax_axis", int)
        )
        y = func.softmax(np.float32([[4, 5, 6]]), axis=-1)
        z = func.add(x, y)
        func.outputs += [func.add(input, z)]

        # Only one node should get folded.
        assert len(func.nodes) == 4
        func.fold_constants().cleanup()
        assert len(func.nodes) == 3

    def test_functions_inside_functions(self, foldable_with_local_functions):
        graph = foldable_with_local_functions
        graph.toposort()
        graph.fold_constants(error_ok=False)
        graph.cleanup()

        assert len(graph.inputs) == 1
        assert len(graph.outputs) == 1
        assert isinstance(graph.outputs[0], Constant)
        assert graph.outputs[0].values == 8

    def test_function_with_unused_input(self, foldable_with_local_functions):
        # Constant folding should still work correctly when a function has unused inputs.
        graph = foldable_with_local_functions

        func_outer = graph.functions[1]
        func_outer.nodes[1].inputs[0] = gs.Constant(
            "Constant_99", values=np.array([3], dtype=np.float32)
        )

        graph.toposort().fold_constants(error_ok=False).cleanup()

        assert len(graph.inputs) == 1
        assert len(graph.outputs) == 1
        assert isinstance(graph.outputs[0], Constant)
        assert graph.outputs[0].values == 12

    def test_inner_function_with_unused_input(self, foldable_with_local_functions):
        graph = foldable_with_local_functions

        func_inner = graph.functions[0]
        func_inner.nodes[0].inputs[0] = gs.Constant(
            "Constant_99", values=np.array([3], dtype=np.float32)
        )

        graph.toposort().fold_constants(error_ok=False).cleanup()

        assert len(graph.inputs) == 1
        assert len(graph.outputs) == 1
        assert isinstance(graph.outputs[0], Constant)
        assert graph.outputs[0].values == 12

    def test_function_with_subgraph(self, foldable_with_local_functions):
        graph = foldable_with_local_functions
        dtype = graph.outputs[0].dtype

        func = Function("func_with_subgraph", ir_version=10)
        func.inputs = [Variable("input")]

        then_graph = Graph(name="Then", ir_version=10)
        then_graph.outputs = [then_graph.identity(func.inputs[0])]

        else_graph = Graph(name="Else", ir_version=10)
        else_graph.functions = graph.functions
        else_graph.outputs = else_graph.FuncInner(
            inputs=[func.inputs[0]], outputs=["else_out"]
        )

        cond = func.less(func.inputs[0], Constant("Zero", np.zeros(1, dtype=dtype)))
        func.outputs = [func.if_op(cond, then_graph, else_graph)]

        graph.functions.append(func)
        graph.outputs = graph.func_with_subgraph(
            inputs=[graph.outputs[0]], outputs=["new_output"]
        )
        graph.toposort().fold_constants(error_ok=False).cleanup()
        assert len(graph.inputs) == 1
        assert len(graph.outputs) == 1
        assert isinstance(graph.outputs[0], Constant)
        assert graph.outputs[0].values == 9

    def test_function_with_complicated_attrs(self):
        # Types of attributes we should test:
        # 1) Has no default value
        # 2) Has default value which is used
        # 3) Has default value which is overridden
        # 4) Confusing attribute name / reference name mappings
        dtype = np.float32
        opset = 18
        func = Function(
            "complicated_func",
            inputs=[Variable("input", dtype=dtype)],
            opset=opset,
            ir_version=10,
        )
        variables = [Variable(f"var{i}", dtype=dtype) for i in range(5)]

        func.nodes.append(
            Node(
                "ConstantOfShape",
                attrs={"value": Node.AttributeRef("ConstantOfShape_value", Tensor)},
                inputs=[func.inputs[0]],
                outputs=[variables[0]],  # shape [2, 3, 4]
            )
        )
        func.nodes.append(
            Node(
                "ReduceSum",
                attrs={"keepdims": Node.AttributeRef("keepdims", int)},
                inputs=[
                    variables[0],
                    Constant("ReduceSum_axis", np.array([1], dtype=int)),
                ],
                outputs=[variables[1]],  # shape [2, 1, 4] when keepDims=True
            )
        )
        func.nodes.append(
            Node(
                "Flatten",
                attrs={"axis": Node.AttributeRef("Flatten_axis", np.int32)},
                inputs=[variables[1]],
                outputs=[variables[2]],  # shape [2, 4] when axis=1
            )
        )
        func.nodes.append(
            Node(
                "Concat",
                attrs={"axis": Node.AttributeRef("axis", np.int64)},
                inputs=[
                    variables[2],
                    Constant("to_concat", values=np.ones((2, 1), dtype=dtype)),
                ],
                outputs=[variables[3]],  # shape [2, 5] when axis=-1
            )
        )
        func.nodes.append(
            Node(
                "Concat",
                attrs={"axis": Node.AttributeRef("Concat_axis", int)},
                inputs=[
                    variables[3],
                    Constant("to_concat_2", values=2 * np.ones((2, 1), dtype=dtype)),
                ],
                outputs=[variables[4]],  # shape [2, 6] when axis=-1
            )
        )
        func.outputs = [variables[4]]

        func.attrs = {
            "ConstantOfShape_value": None,  # no default value
            "axis": None,  # no default value
            "Flatten_axis": 0,
            "keepdims": True,
            "Concat_axis": -1,
        }
        graph = Graph(opset=opset, functions=[func], ir_version=10)
        input = Constant("shape", values=np.array([2, 3, 4], dtype=int))
        output = graph.complicated_func(
            inputs=[input],
            outputs=["output"],
            attrs={
                "ConstantOfShape_value": Constant(
                    "three", values=np.array([3], dtype=dtype)
                ),
                "axis": -1,
                "Flatten_axis": 2,  # overrides default value
            },
        )[0]
        output.dtype = dtype
        graph.inputs = [input]
        graph.outputs = [output]
        graph.fold_constants(error_ok=False).cleanup()

        assert len(graph.inputs) == 1
        assert len(graph.outputs) == 1
        assert isinstance(graph.outputs[0], Constant)
        assert np.all(
            graph.outputs[0].values
            == np.array([[9, 9, 9, 9, 1, 2], [9, 9, 9, 9, 1, 2]])
        )


class TestIO(object):
    def test_io_cannot_be_sync_list_on_init(self):
        inp = Variable("input0", shape=(1, 3), dtype=np.float32)
        out = Variable("input1", shape=(1, 3), dtype=np.float32)

        node = Node("Add", inputs=[inp], outputs=[out])
        assert isinstance(node.inputs, SynchronizedList)
        assert isinstance(node.outputs, SynchronizedList)

        graph = Graph(nodes=[node], inputs=node.inputs, outputs=node.outputs)
        assert not isinstance(graph.inputs, SynchronizedList)
        assert not isinstance(graph.outputs, SynchronizedList)

    def test_io_cannot_be_sync_list_on_assign(self):
        inp = Variable("input0", shape=(1, 3), dtype=np.float32)
        out = Variable("input1", shape=(1, 3), dtype=np.float32)

        node = Node("Add", inputs=[inp], outputs=[out])
        assert isinstance(node.inputs, SynchronizedList)
        assert isinstance(node.outputs, SynchronizedList)

        graph = Graph(nodes=[node], inputs=[], outputs=[])
        graph.inputs = node.inputs
        graph.outputs = node.outputs

        assert not isinstance(graph.inputs, SynchronizedList)
        assert not isinstance(graph.outputs, SynchronizedList)


class TestFunctionList(object):
    def test_shared_function_list(self):
        subgraph = Graph()
        node = Node("Test", attrs={"test_attr": subgraph})
        main_graph = Graph(nodes=[node])

        new_func = Function("func")
        main_graph.functions.append(new_func)
        assert len(main_graph.functions) == 1
        assert len(subgraph.functions) == 1
        assert subgraph.functions[0] == new_func

    def test_set_function_list(self):
        subgraph = Graph()
        node = Node("Test", attrs={"test_attr": subgraph})
        main_graph = Graph(nodes=[node])

        new_func = Function("func")
        main_graph.functions = [new_func]
        assert len(main_graph.functions) == 1
        assert len(subgraph.functions) == 1
        assert subgraph.functions[0] == new_func

    def test_merge_funcs(self):
        func1 = Function("func1", domain="domain1")
        func2 = Function("func1", domain="domain2")
        func3 = Function("func2", domain="domain1")
        func4 = Function("func2", domain="domain2")
        all_funcs = [func1, func2, func3, func4]

        subgraph1 = Graph(functions=[func1, func2])

        node1 = Node("Test", attrs={"test_attr": subgraph1})
        subgraph2 = Graph(functions=[func2, func3], nodes=[node1])

        node2 = Node("Test", attrs={"test_attr": subgraph2})
        main_graph = Graph(functions=[func3, func4], nodes=[node2])

        assert len(main_graph.functions) == len(all_funcs)
        assert len(subgraph1.functions) == len(all_funcs)
        assert len(subgraph2.functions) == len(all_funcs)

        new_func = Function("func3")
        main_graph.functions.append(new_func)
        assert new_func in subgraph1.functions
        assert new_func in subgraph2.functions