import unittest

import numpy
import onnx_light.onnx as onnxl
import onnx_light.onnx.helper as oh
import onnx_light.onnx.numpy_helper
from onnx_light.ext_test_case import ExtTestCase
from onnx_light.onnx_core import shape_inference as si


class TestInPlaceReuse(ExtTestCase):
    """Python tests for the ``compute_inplace_reuse`` binding exposed by
    ``onnx_light.onnx_py._onnxpycore.shape_inference``."""

    _LONG_SYMBOLIC_DIM_NAME = "very_long_dynamic_dimension_name_" * 16 + "N"

    @staticmethod
    def _reuse_pairs(reuse):
        """Normalises the nested ``InPlaceReuse`` result into a list of
        lists of ``(output_index, input_index)`` tuples."""
        return [[(r.output_index, r.input_index) for r in node] for node in reuse]

    def test_inplace_reuse_struct(self):
        r = si.InPlaceReuse()
        r.output_index = 1
        r.input_index = 2
        self.assertEqual(r.output_index, 1)
        self.assertEqual(r.input_index, 2)
        # The default kind is kEqual (same-sized buffer).
        self.assertEqual(r.kind, si.InPlaceReuseKind.kEqual)
        self.assertEqual(repr(r), "InPlaceReuse(output_index=1, input_index=2, kind=kEqual)")
        r.kind = si.InPlaceReuseKind.kGreater
        self.assertEqual(r.kind, si.InPlaceReuseKind.kGreater)
        self.assertEqual(repr(r), "InPlaceReuse(output_index=1, input_index=2, kind=kGreater)")

    def _build_model(self, nodes, inputs, outputs):
        graph = oh.make_graph(nodes, "g", inputs, outputs)
        model = oh.make_model(graph, opset_imports=[oh.make_opsetid("", 18)])
        model.ir_version = 8
        return model

    def test_abs_chain_reuses_intermediates(self):
        nodes = [
            oh.make_node("Abs", ["X"], ["A"]),
            oh.make_node("Abs", ["A"], ["B"]),
            oh.make_node("Abs", ["B"], ["Y"]),
        ]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [3, 4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [3, 4])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)
        reuse = self._reuse_pairs(si.compute_inplace_reuse(ctx, model.graph))

        # Node 0 reads the declared graph input X (must not be overwritten).
        self.assertEqual(reuse, [[], [(0, 0)], [(0, 0)]])
        # The reused buffers are the same size as the outputs.
        raw = si.compute_inplace_reuse(ctx, model.graph)
        self.assertEqual(raw[1][0].kind, si.InPlaceReuseKind.kEqual)
        self.assertEqual(raw[2][0].kind, si.InPlaceReuseKind.kEqual)

    def test_larger_input_buffer_reported_as_greater(self):
        # A is INT64 (8 bytes/elem); Y is INT32 (4 bytes/elem). A's buffer is
        # strictly larger than Y, so Y may still reuse it in place.
        nodes = [
            oh.make_node("Cast", ["X"], ["A"], to=onnxl.TensorProto.INT64),
            oh.make_node("Cast", ["A"], ["Y"], to=onnxl.TensorProto.INT32),
        ]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.INT32, [4])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)
        raw = si.compute_inplace_reuse(ctx, model.graph)
        reuse = self._reuse_pairs(raw)

        self.assertEqual(reuse, [[], [(0, 0)]])
        self.assertEqual(raw[1][0].kind, si.InPlaceReuseKind.kGreater)

    def test_value_read_twice_reused_only_at_last_use(self):
        nodes = [
            oh.make_node("Abs", ["X"], ["A"]),
            oh.make_node("Abs", ["A"], ["B"]),
            oh.make_node("Add", ["A", "B"], ["Y"]),
        ]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [3, 4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [3, 4])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)
        reuse = self._reuse_pairs(si.compute_inplace_reuse(ctx, model.graph))

        # A is alive at node 2, so node 1 cannot reuse it; node 2 reuses A.
        self.assertEqual(reuse, [[], [], [(0, 0)]])

    def test_transpose_same_byte_size_reported_as_equal(self):
        nodes = [
            oh.make_node("Abs", ["X"], ["A"]),
            oh.make_node("Transpose", ["A"], ["B"]),
            oh.make_node("Transpose", ["B"], ["Y"]),
        ]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [3, 4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [3, 4])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)
        raw = si.compute_inplace_reuse(ctx, model.graph)
        reuse = self._reuse_pairs(raw)

        # Both transposes change layout but keep the same byte size.
        self.assertEqual(reuse, [[], [(0, 0)], [(0, 0)]])
        self.assertEqual(raw[1][0].kind, si.InPlaceReuseKind.kEqual)
        self.assertEqual(raw[2][0].kind, si.InPlaceReuseKind.kEqual)

    def test_transpose_symbolic_dim_reported_as_equal(self):
        """Tests that Transpose with a symbolic batch dimension reports kEqual.

        When the input shape contains a symbolic dimension (e.g. batch), the
        byte sizes cannot be compared concretely.  The fix uses the symbolic
        ByteSizeExpr: both [batch, 4, 4] and [4, batch, 4] simplify to the same
        canonical expression, so kEqual is correctly returned.
        """
        nodes = [oh.make_node("Abs", ["X"], ["A"]), oh.make_node("Transpose", ["A"], ["Y"])]
        # Build a minimal model proto; shapes are injected via ctx.set below.
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, None)
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, None)
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        ctx.set("X", si.SymTensor(onnxl.TensorProto.FLOAT, ["batch", 4, 4]))
        ctx.set("A", si.SymTensor(onnxl.TensorProto.FLOAT, ["batch", 4, 4]))
        ctx.set("Y", si.SymTensor(onnxl.TensorProto.FLOAT, [4, "batch", 4]))

        raw = si.compute_inplace_reuse(ctx, model.graph)
        reuse = self._reuse_pairs(raw)

        # Node 0: Abs(X) → A — X is a declared graph input, must not be reused.
        # Node 1: Transpose(A) → Y — A's byte size equals Y's, so kEqual.
        self.assertEqual(reuse, [[], [(0, 0)]])
        self.assertEqual(raw[1][0].kind, si.InPlaceReuseKind.kEqual)

    def test_transpose_square_shape_reported_as_equal(self):
        """Tests that a square-matrix Transpose is reported as kEqual.

        For a square-matrix Transpose the output shape equals the input shape
        ([4, 4] -> [4, 4]).  SameStorage returns true, so the inplace kind is kEqual.
        """
        nodes = [oh.make_node("Abs", ["X"], ["A"]), oh.make_node("Transpose", ["A"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [4, 4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [4, 4])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)
        raw = si.compute_inplace_reuse(ctx, model.graph)
        reuse = self._reuse_pairs(raw)

        self.assertEqual(reuse, [[], [(0, 0)]])
        self.assertEqual(raw[1][0].kind, si.InPlaceReuseKind.kEqual)

    def test_unsqueeze_dynamic_axes_equal_reuse(self):
        nodes = [
            oh.make_node("Abs", ["X"], ["A"]),
            oh.make_node("Unsqueeze", ["A", "axes"], ["Y"]),
        ]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, None)
        axes = oh.make_tensor_value_info("axes", onnxl.TensorProto.INT64, None)
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, None)
        model = self._build_model(nodes, [x, axes], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)
        raw = si.compute_inplace_reuse(ctx, model.graph)
        reuse = self._reuse_pairs(raw)

        # Unsqueeze keeps the same underlying buffer as its data input even when
        # the axes value is dynamic and shape inference cannot prove byte sizes.
        self.assertEqual(reuse, [[], [(0, 0)]])
        self.assertEqual(raw[1][0].kind, si.InPlaceReuseKind.kEqual)

    def test_squeeze_dynamic_axes_equal_reuse(self):
        nodes = [oh.make_node("Abs", ["X"], ["A"]), oh.make_node("Squeeze", ["A", "axes"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, None)
        axes = oh.make_tensor_value_info("axes", onnxl.TensorProto.INT64, None)
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, None)
        model = self._build_model(nodes, [x, axes], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)
        raw = si.compute_inplace_reuse(ctx, model.graph)
        reuse = self._reuse_pairs(raw)

        # Squeeze keeps the same underlying buffer as its data input even when
        # axes is dynamic and shape inference cannot prove byte sizes.
        self.assertEqual(reuse, [[], [(0, 0)]])
        self.assertEqual(raw[1][0].kind, si.InPlaceReuseKind.kEqual)

    def test_unsqueeze_is_not_reused_before_last_use(self):
        nodes = [
            oh.make_node("Abs", ["X"], ["A"]),
            oh.make_node("Unsqueeze", ["A", "axes"], ["Y"]),
            oh.make_node("Abs", ["A"], ["Z"]),
        ]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, None)
        axes = oh.make_tensor_value_info("axes", onnxl.TensorProto.INT64, None)
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, None)
        z = oh.make_tensor_value_info("Z", onnxl.TensorProto.FLOAT, None)
        model = self._build_model(nodes, [x, axes], [y, z])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)
        reuse = self._reuse_pairs(si.compute_inplace_reuse(ctx, model.graph))

        # A is consumed again by node 2, so node 1 cannot alias it.
        self.assertEqual(reuse, [[], [], [(0, 0)]])

    def test_graph_output_input_is_not_reused(self):
        nodes = [oh.make_node("Abs", ["X"], ["A"]), oh.make_node("Abs", ["A"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [2, 2])
        a = oh.make_tensor_value_info("A", onnxl.TensorProto.FLOAT, [2, 2])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [2, 2])
        model = self._build_model(nodes, [x], [a, y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)
        reuse = self._reuse_pairs(si.compute_inplace_reuse(ctx, model.graph))

        # A is a declared graph output; node 1 must not overwrite it.
        self.assertEqual(reuse, [[], []])

    @staticmethod
    def _node_metadata(node):
        return {m.key: m.value for m in node.metadata_props}

    def test_write_inplace_reuse_to_metadata(self):
        nodes = [
            oh.make_node("Abs", ["X"], ["A"]),
            oh.make_node("Abs", ["A"], ["B"]),
            oh.make_node("Abs", ["B"], ["Y"]),
        ]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [3, 4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [3, 4])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)
        si.write_inplace_reuse_to_metadata(ctx, model.graph)

        # Node 0 reads the declared graph input X for the last time.
        self.assertEqual(
            self._node_metadata(model.graph.node[0]), {"onnx_light.not_used_after": "X"}
        )
        self.assertEqual(
            self._node_metadata(model.graph.node[1]),
            {"onnx_light.inplace_reuse": "0:0:equal", "onnx_light.release_after": "A"},
        )
        self.assertEqual(
            self._node_metadata(model.graph.node[2]),
            {"onnx_light.inplace_reuse": "0:0:equal", "onnx_light.release_after": "B"},
        )

    def test_write_inplace_reuse_to_metadata_greater(self):
        nodes = [
            oh.make_node("Cast", ["X"], ["A"], to=onnxl.TensorProto.INT64),
            oh.make_node("Cast", ["A"], ["Y"], to=onnxl.TensorProto.INT32),
        ]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.INT32, [4])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)
        si.write_inplace_reuse_to_metadata(ctx, model.graph)

        self.assertEqual(
            self._node_metadata(model.graph.node[0]), {"onnx_light.not_used_after": "X"}
        )
        self.assertEqual(
            self._node_metadata(model.graph.node[1]),
            {"onnx_light.inplace_reuse": "0:0:greater", "onnx_light.release_after": "A"},
        )

    def test_write_inplace_reuse_to_metadata_updates_existing_key(self):
        nodes = [oh.make_node("Abs", ["X"], ["A"]), oh.make_node("Abs", ["A"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [3, 4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [3, 4])
        model = self._build_model(nodes, [x], [y])
        # Pre-existing entry under the same key should be replaced in place.
        model.graph.node[1].add_metadata("onnx_light.inplace_reuse", "stale")

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)
        si.write_inplace_reuse_to_metadata(ctx, model.graph)

        self.assertEqual(
            self._node_metadata(model.graph.node[0]), {"onnx_light.not_used_after": "X"}
        )
        self.assertEqual(
            self._node_metadata(model.graph.node[1]),
            {"onnx_light.inplace_reuse": "0:0:equal", "onnx_light.release_after": "A"},
        )

    def test_write_inplace_reuse_to_metadata_tracks_initializer_last_use(self):
        nodes = [oh.make_node("Add", ["X", "W"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [2, 2])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [2, 2])
        w = onnx_light.onnx.numpy_helper.from_array(
            numpy.ones((2, 2), dtype=numpy.float32), name="W"
        )
        model = self._build_model(nodes, [x], [y])
        model.graph.initializer.extend([w])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)
        si.write_inplace_reuse_to_metadata(ctx, model.graph)

        self.assertEqual(
            self._node_metadata(model.graph.node[0]), {"onnx_light.not_used_after": "X;W"}
        )

    def test_allow_input_overwrite_reuses_graph_input(self):
        nodes = [oh.make_node("Abs", ["X"], ["A"]), oh.make_node("Abs", ["A"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [3, 4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [3, 4])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)

        # By default the graph input X must not be overwritten.
        reuse = self._reuse_pairs(si.compute_inplace_reuse(ctx, model.graph))
        self.assertEqual(reuse, [[], [(0, 0)]])

        # Opt-in: node 0 is the only use of X, so it may reuse it in place.
        reuse_ovw = self._reuse_pairs(
            si.compute_inplace_reuse(ctx, model.graph, allow_input_overwrite=True)
        )
        self.assertEqual(reuse_ovw, [[(0, 0)], [(0, 0)]])

    def test_allow_input_overwrite_keeps_input_that_is_also_output(self):
        nodes = [oh.make_node("Abs", ["X"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [2, 2])
        x_out = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [2, 2])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [2, 2])
        model = self._build_model(nodes, [x], [x_out, y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)

        # X is also a declared output, so it stays protected even with the option.
        reuse = self._reuse_pairs(
            si.compute_inplace_reuse(ctx, model.graph, allow_input_overwrite=True)
        )
        self.assertEqual(reuse, [[]])

    def test_inplace_context_stores_result(self):
        nodes = [
            oh.make_node("Abs", ["X"], ["A"]),
            oh.make_node("Abs", ["A"], ["B"]),
            oh.make_node("Abs", ["B"], ["Y"]),
        ]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [3, 4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [3, 4])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)

        inplace = si.ComputeContext()
        self.assertEqual(len(inplace), 0)
        inplace.compute_inplace_reuse_graph(model.graph, ctx)

        self.assertEqual(len(inplace), 3)
        # The stored result matches the free-function wrapper.
        self.assertEqual(self._reuse_pairs(inplace.reuse), [[], [(0, 0)], [(0, 0)]])
        self.assertEqual(self._reuse_pairs([inplace.node_reuse(1)]), [[(0, 0)]])
        self.assertEqual(inplace.node_reuse(1)[0].kind, si.InPlaceReuseKind.kEqual)
        # node_reuse rejects an out-of-bounds index.
        with self.assertRaises(IndexError):
            inplace.node_reuse(3)
        # clear empties the stored result.
        inplace.clear()
        self.assertEqual(len(inplace), 0)

    def test_inplace_context_allow_input_overwrite(self):
        nodes = [oh.make_node("Abs", ["X"], ["A"]), oh.make_node("Abs", ["A"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [3, 4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [3, 4])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)

        inplace = si.ComputeContext()
        inplace.compute_inplace_reuse_graph(model.graph, ctx, allow_input_overwrite=True)
        self.assertEqual(self._reuse_pairs(inplace.reuse), [[(0, 0)], [(0, 0)]])

    def test_inplace_context_memory_tracks_sources_and_reuse(self):
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.INT32, [4])
        shape_init = oh.make_tensor("S", onnxl.TensorProto.INT64, [1], [4])
        graph = oh.make_graph(
            [
                oh.make_node("P", ["X", "S"], ["A"]),
                oh.make_node("Q", ["A"], ["B"]),
                oh.make_node("R", ["B"], ["Y"]),
            ],
            "g",
            [x],
            [y],
            [shape_init],
        )
        model = oh.make_model(graph, opset_imports=[oh.make_opsetid("", 18)])
        model.ir_version = 8

        ctx = si.ShapesContext()
        ctx.set("X", si.SymTensor(onnxl.TensorProto.FLOAT, [4]))
        ctx.set("S", si.SymTensor(onnxl.TensorProto.INT64, [1]))
        ctx.set("A", si.SymTensor(onnxl.TensorProto.FLOAT, [4]))
        ctx.set("B", si.SymTensor(onnxl.TensorProto.INT64, [4]))
        ctx.set("Y", si.SymTensor(onnxl.TensorProto.INT32, [4]))

        inplace = si.ComputeContext()
        inplace.compute_inplace_reuse_graph(model.graph, ctx, value_tags={"S": "shape"})

        self.assertEqual(self._reuse_pairs(inplace.reuse), [[], [], [(0, 0)]])
        self.assertEqual(len(inplace.memory), 3)

        mem0 = inplace.node_memory(0)
        self.assertEqual(
            list(mem0.keys()),
            [
                "already_allocated_bytes",
                "initializers",
                "inputs",
                "intermediates",
                "output_allocation_bytes",
                "outputs",
                "total_bytes",
            ],
        )
        self.assertEqual(mem0["total_bytes"], 40)
        self.assertEqual(mem0["inputs"], {"": 16})
        self.assertEqual(mem0["already_allocated_bytes"], 24)
        self.assertEqual(mem0["output_allocation_bytes"], 16)
        self.assertEqual(mem0["inputs"], {"": 16})
        self.assertEqual(mem0["initializers"], {"shape": 8})
        self.assertEqual(mem0["intermediates"], {})
        self.assertEqual(mem0["outputs"], {"": 16})

        mem1 = inplace.node_memory(1)
        self.assertEqual(mem1["total_bytes"], 72)
        self.assertEqual(mem1["already_allocated_bytes"], 40)
        self.assertEqual(mem1["output_allocation_bytes"], 32)
        self.assertEqual(mem1["inputs"], {"": 16})
        self.assertEqual(mem1["initializers"], {"shape": 8})
        self.assertEqual(mem1["intermediates"], {"": 16})
        self.assertEqual(mem1["outputs"], {"": 32})

        mem2 = inplace.node_memory(2)
        self.assertEqual(mem2["total_bytes"], 56)
        self.assertEqual(mem2["already_allocated_bytes"], 56)
        self.assertEqual(mem2["output_allocation_bytes"], 0)
        self.assertEqual(mem2["inputs"], {"": 16})
        self.assertEqual(mem2["initializers"], {"shape": 8})
        self.assertEqual(mem2["intermediates"], {"": 32})
        self.assertEqual(mem2["outputs"], {})

        with self.assertRaises(IndexError):
            inplace.node_memory(3)

    def test_inplace_context_memory_keeps_symbolic_shapes(self):
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, ["N"])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.INT32, ["N"])
        shape_init = oh.make_tensor("S", onnxl.TensorProto.INT64, [1], [4])
        graph = oh.make_graph(
            [oh.make_node("P", ["X", "S"], ["A"]), oh.make_node("Q", ["A"], ["Y"])],
            "g",
            [x],
            [y],
            [shape_init],
        )
        model = oh.make_model(graph, opset_imports=[oh.make_opsetid("", 18)])
        model.ir_version = 8

        ctx = si.ShapesContext()
        ctx.set("X", si.SymTensor(onnxl.TensorProto.FLOAT, ["N"]))
        ctx.set("S", si.SymTensor(onnxl.TensorProto.INT64, [1]))
        ctx.set("A", si.SymTensor(onnxl.TensorProto.FLOAT, ["N"]))
        ctx.set("Y", si.SymTensor(onnxl.TensorProto.INT32, ["N"]))

        inplace = si.ComputeContext()
        inplace.compute_inplace_reuse_graph(model.graph, ctx, value_tags={"S": "shape"})

        mem0 = inplace.node_memory(0)
        self.assertEqual(mem0["outputs"], {"": "4*N"})
        self.assertEqual(mem0["total_bytes"], "8*N+8")
        self.assertEqual(mem0["already_allocated_bytes"], "4*N+8")
        self.assertEqual(mem0["output_allocation_bytes"], "4*N")
        self.assertEqual(mem0["inputs"], {"": "4*N"})
        self.assertEqual(mem0["initializers"], {"shape": 8})
        self.assertEqual(mem0["intermediates"], {})
        self.assertEqual(mem0["outputs"], {"": "4*N"})

        mem1 = inplace.node_memory(1)
        # A (FLOAT [N]) and Y (INT32 [N]) have equal byte sizes (4*N each), so A's
        # buffer is reused for Y — no extra output allocation for this node.
        self.assertEqual(mem1["total_bytes"], "8*N+8")
        self.assertEqual(mem1["already_allocated_bytes"], "8*N+8")
        self.assertEqual(mem1["output_allocation_bytes"], 0)
        self.assertEqual(mem1["inputs"], {"": "4*N"})
        self.assertEqual(mem1["initializers"], {"shape": 8})
        self.assertEqual(mem1["intermediates"], {"": "4*N"})
        self.assertEqual(mem1["outputs"], {})

    def test_inplace_context_memory_simplifies_repeated_symbolic_sums(self):
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, ["N"])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, ["N"])
        nodes = [oh.make_node("P", ["X"], [f"A{i}"]) for i in range(5)]
        nodes.append(oh.make_node("Q", [f"A{i}" for i in range(5)], ["Y"]))
        graph = oh.make_graph(nodes, "g", [x], [y])
        model = oh.make_model(graph, opset_imports=[oh.make_opsetid("", 18)])
        model.ir_version = 8

        ctx = si.ShapesContext()
        ctx.set("X", si.SymTensor(onnxl.TensorProto.FLOAT, ["N"]))
        for i in range(5):
            ctx.set(f"A{i}", si.SymTensor(onnxl.TensorProto.FLOAT, ["N"]))
        ctx.set("Y", si.SymTensor(onnxl.TensorProto.FLOAT, ["N"]))

        inplace = si.ComputeContext()
        inplace.compute_inplace_reuse_graph(model.graph, ctx)

        mem4 = inplace.node_memory(4)
        self.assertEqual(mem4["already_allocated_bytes"], "20*N")
        self.assertEqual(mem4["output_allocation_bytes"], "4*N")
        self.assertEqual(mem4["total_bytes"], "24*N")
        self.assertEqual(mem4["inputs"], {"": "4*N"})
        self.assertEqual(mem4["intermediates"], {"": "16*N"})
        self.assertEqual(mem4["outputs"], {"": "4*N"})

        mem5 = inplace.node_memory(5)
        self.assertEqual(mem5["already_allocated_bytes"], "24*N")
        self.assertEqual(mem5["output_allocation_bytes"], 0)
        self.assertEqual(mem5["total_bytes"], "24*N")
        self.assertEqual(mem5["inputs"], {"": "4*N"})
        self.assertEqual(mem5["intermediates"], {"": "20*N"})
        self.assertEqual(mem5["outputs"], {})

    def test_if_subgraph_local_shadowing_excluded_from_captures(self):
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [3, 4])
        cond = oh.make_tensor_value_info("cond", onnxl.TensorProto.BOOL, [])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [3, 4])

        then_branch = oh.make_graph(
            [oh.make_node("Abs", ["X"], ["A"]), oh.make_node("Identity", ["A"], ["T"])],
            "then",
            [],
            [oh.make_tensor_value_info("T", onnxl.TensorProto.FLOAT, [3, 4])],
        )
        else_branch = oh.make_graph(
            [oh.make_node("Abs", ["X"], ["A"]), oh.make_node("Identity", ["A"], ["E"])],
            "else",
            [],
            [oh.make_tensor_value_info("E", onnxl.TensorProto.FLOAT, [3, 4])],
        )

        model = self._build_model(
            [
                oh.make_node("Abs", ["X"], ["A"]),
                oh.make_node(
                    "If", ["cond"], ["B"], then_branch=then_branch, else_branch=else_branch
                ),
                oh.make_node("Abs", ["B"], ["Y"]),
            ],
            [x, cond],
            [y],
        )

        ctx = si.ShapesContext()
        ctx.set("X", si.SymTensor(onnxl.TensorProto.FLOAT, [3, 4]))
        ctx.set("A", si.SymTensor(onnxl.TensorProto.FLOAT, [3, 4]))
        ctx.set("B", si.SymTensor(onnxl.TensorProto.FLOAT, [3, 4]))
        ctx.set("Y", si.SymTensor(onnxl.TensorProto.FLOAT, [3, 4]))

        inplace = si.ComputeContext()
        inplace.compute_inplace_reuse_graph(model.graph, ctx)

        mem1 = inplace.node_memory(1)
        self.assertEqual(mem1["already_allocated_bytes"], 48)
        self.assertEqual(mem1["inputs"], {"": 48})
        self.assertEqual(mem1["intermediates"], {})
        self.assertEqual(mem1["output_allocation_bytes"], 48)
        self.assertEqual(mem1["outputs"], {"": 48})
        self.assertEqual(mem1["total_bytes"], 96)

    def test_inplace_context_write_to_metadata(self):
        nodes = [
            oh.make_node("Abs", ["X"], ["A"]),
            oh.make_node("Abs", ["A"], ["B"]),
            oh.make_node("Abs", ["B"], ["Y"]),
        ]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [3, 4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [3, 4])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)

        inplace = si.ComputeContext()
        inplace.compute_inplace_reuse_graph(model.graph, ctx)
        inplace.write_to_metadata(model.graph)

        self.assertEqual(
            self._node_metadata(model.graph.node[0]), {"onnx_light.not_used_after": "X"}
        )
        self.assertEqual(
            self._node_metadata(model.graph.node[1]),
            {"onnx_light.inplace_reuse": "0:0:equal", "onnx_light.release_after": "A"},
        )
        self.assertEqual(
            self._node_metadata(model.graph.node[2]),
            {"onnx_light.inplace_reuse": "0:0:equal", "onnx_light.release_after": "B"},
        )

    def test_shape_tag_release_info_populates_shape_tagged(self):
        # Shape(X)->S, Reshape(X, S)->Y: S is shape-tagged and released after Reshape.
        nodes = [oh.make_node("Shape", ["X"], ["S"]), oh.make_node("Reshape", ["X", "S"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [2, 3])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [2, 3])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)

        value_tags = {"S": "shape"}
        inplace = si.ComputeContext()
        inplace.compute_inplace_reuse_graph(model.graph, ctx, value_tags=value_tags)

        self.assertEqual(len(inplace), 2)
        # S is released after the Reshape node (index 1).
        self.assertEqual(inplace.release_after_shape_tagged[0], [])
        self.assertEqual(inplace.release_after_shape_tagged[1], ["S"])
        # Per-node accessor mirrors the list.
        self.assertEqual(inplace.node_release_after_shape_tagged(1), ["S"])
        with self.assertRaises(IndexError):
            inplace.node_release_after_shape_tagged(2)

    def test_compute_value_and_node_tags_method(self):
        nodes = [oh.make_node("Shape", ["X"], ["S"]), oh.make_node("Reshape", ["X", "S"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [2, 3])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [2, 3])
        model = self._build_model(nodes, [x], [y])

        compute = si.ComputeContext()
        value_tags, node_tags = compute.compute_value_and_node_tags(model.graph)

        self.assertEqual(value_tags["S"], "shape")
        self.assertEqual(node_tags[0], "shape")
        self.assertEqual(len(node_tags), 2)
        self.assertEqual(compute.value_tags, value_tags)
        self.assertEqual(compute.node_tags, node_tags)
        self.assertEqual(compute.node_tag(0), "shape")
        self.assertEqual(compute.node_tag(1), node_tags[1])
        with self.assertRaises(IndexError):
            compute.node_tag(2)

    def test_custom_value_tag_function_for_custom_op(self):
        nodes = [oh.make_node("CustomIdentity", ["X"], ["Y"], domain="com.acme")]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [2, 3])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.INT64, [2])
        model = self._build_model(nodes, [x], [y])

        compute = si.ComputeContext()
        value_tags, node_tags = compute.compute_value_and_node_tags(model.graph)
        self.assertEqual(value_tags["Y"], "weight")
        self.assertEqual(node_tags[0], "weight")

        calls = []

        def _custom_shape_tag(ctx, node, node_index):
            calls.append((str(node.domain), str(node.op_type), int(node_index)))
            ctx.set_node_tag(node_index, "shape")
            ctx.try_set_value_tag(str(node.output[0]), "shape")

        self.assertFalse(compute.has_custom_value_tag_function("com.acme", "CustomIdentity"))
        compute.set_custom_value_tag_function("com.acme", "CustomIdentity", _custom_shape_tag)
        self.assertTrue(compute.has_custom_value_tag_function("com.acme", "CustomIdentity"))
        self.assertIn("com.acme:CustomIdentity", list(compute.custom_value_tag_keys()))

        value_tags, node_tags = compute.compute_value_and_node_tags(model.graph)
        self.assertEqual(value_tags["Y"], "shape")
        self.assertEqual(node_tags[0], "shape")
        self.assertGreaterEqual(len(calls), 1)
        self.assertEqual(calls[0], ("com.acme", "CustomIdentity", 0))

        self.assertTrue(compute.remove_custom_value_tag_function("com.acme", "CustomIdentity"))
        self.assertFalse(compute.remove_custom_value_tag_function("com.acme", "CustomIdentity"))
        compute.clear_custom_value_tag_functions()
        self.assertEqual(list(compute.custom_value_tag_keys()), [])

    def test_compute_context_accepts_verbose_parameter(self):
        nodes = [oh.make_node("Shape", ["X"], ["S"]), oh.make_node("Reshape", ["X", "S"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [2, 3])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [2, 3])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)

        compute = si.ComputeContext()
        value_tags, _ = compute.compute_value_and_node_tags(model.graph, verbose=1)
        compute.compute_inplace_reuse_graph(model.graph, ctx, value_tags=value_tags)
        self.assertEqual(compute.compute_release_after_shape_tagged(), [[], ["S"]])

    def test_infer_value_and_node_tags_binding_removed(self):
        self.assertFalse(hasattr(si._C, "infer_value_and_node_tags"))
        self.assertIs(si.infer_value_and_node_tags, si.compute_value_and_node_tags)

    def test_shape_tag_release_info_uses_stored_tags(self):
        nodes = [oh.make_node("Shape", ["X"], ["S"]), oh.make_node("Reshape", ["X", "S"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [2, 3])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [2, 3])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)

        inplace = si.ComputeContext()
        inplace.compute_value_and_node_tags(model.graph)
        inplace.compute_inplace_reuse_graph(model.graph, ctx)

        self.assertEqual(inplace.release_after_shape_tagged[0], [])
        self.assertEqual(inplace.release_after_shape_tagged[1], ["S"])

    def test_shape_tag_release_info_writes_metadata(self):
        nodes = [oh.make_node("Shape", ["X"], ["S"]), oh.make_node("Reshape", ["X", "S"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [2, 3])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [2, 3])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)

        value_tags = {"S": "shape"}
        inplace = si.ComputeContext()
        inplace.compute_inplace_reuse_graph(model.graph, ctx, value_tags=value_tags)
        inplace.write_to_metadata(model.graph)

        # Node 0 (Shape) has no release metadata because X is still needed by node 1.
        self.assertEqual(self._node_metadata(model.graph.node[0]), {})
        meta1 = self._node_metadata(model.graph.node[1])
        # kReleaseAfterMetadataKey lists S.
        self.assertEqual(meta1.get("onnx_light.release_after"), "S")
        self.assertEqual(meta1.get("onnx_light.not_used_after"), "X")
        # kReleaseAfterShapeTagMetadataKey also lists S (it is shape-tagged).
        self.assertEqual(meta1.get("onnx_light.release_after_shape_tag"), "S")

    def test_shape_tag_release_info_with_long_symbolic_names(self):
        nodes = [oh.make_node("Shape", ["X"], ["S"]), oh.make_node("Reshape", ["X", "S"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, None)
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, None)
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        # Keeps the symbolic byte-size expression intentionally large so the test
        # exercises the caching path added for expensive simplification work.
        ctx.set("X", si.SymTensor(onnxl.TensorProto.FLOAT, [self._LONG_SYMBOLIC_DIM_NAME]))
        ctx.set("S", si.SymTensor(onnxl.TensorProto.INT64, [1]))
        ctx.set("Y", si.SymTensor(onnxl.TensorProto.FLOAT, [self._LONG_SYMBOLIC_DIM_NAME]))

        inplace = si.ComputeContext()
        inplace.compute_inplace_reuse_graph(model.graph, ctx, value_tags={"S": "shape"})

        self.assertEqual(inplace.release_after_shape_tagged, [[], ["S"]])
        mem0 = inplace.node_memory(0)
        self.assertEqual(mem0["inputs"], {"": f"4*{self._LONG_SYMBOLIC_DIM_NAME}"})
        self.assertEqual(mem0["output_allocation_bytes"], 8)
        self.assertEqual(mem0["total_bytes"], f"4*{self._LONG_SYMBOLIC_DIM_NAME}+8")

        inplace.write_to_metadata(model.graph)
        meta1 = self._node_metadata(model.graph.node[1])
        self.assertEqual(meta1.get("onnx_light.release_after"), "S")
        self.assertEqual(meta1.get("onnx_light.not_used_after"), "X")
        self.assertEqual(meta1.get("onnx_light.release_after_shape_tag"), "S")

    def test_shape_tag_no_value_tags_leaves_shape_tagged_empty(self):
        nodes = [oh.make_node("Abs", ["X"], ["A"]), oh.make_node("Abs", ["A"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [3, 4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [3, 4])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)

        inplace = si.ComputeContext()
        inplace.compute_inplace_reuse_graph(model.graph, ctx)
        inplace.write_to_metadata(model.graph)

        # Without value_tags the shape-tagged vector is itself empty.
        self.assertEqual(inplace.release_after_shape_tagged, [])
        # kReleaseAfterShapeTagMetadataKey must not appear.
        for node in model.graph.node:
            self.assertNotIn("onnx_light.release_after_shape_tag", self._node_metadata(node))

    def test_write_inplace_reuse_to_metadata_with_value_tags(self):
        nodes = [oh.make_node("Shape", ["X"], ["S"]), oh.make_node("Reshape", ["X", "S"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [2, 3])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [2, 3])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)

        value_tags = {"S": "shape"}
        si.write_inplace_reuse_to_metadata(ctx, model.graph, value_tags=value_tags)

        meta1 = self._node_metadata(model.graph.node[1])
        self.assertEqual(meta1.get("onnx_light.release_after"), "S")
        self.assertEqual(meta1.get("onnx_light.not_used_after"), "X")
        self.assertEqual(meta1.get("onnx_light.release_after_shape_tag"), "S")

    def test_shape_tag_clear_resets_shape_tagged(self):
        nodes = [oh.make_node("Shape", ["X"], ["S"]), oh.make_node("Reshape", ["X", "S"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [2, 3])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [2, 3])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)

        inplace = si.ComputeContext()
        inplace.compute_inplace_reuse_graph(model.graph, ctx, value_tags={"S": "shape"})
        self.assertEqual(len(inplace), 2)
        inplace.clear()
        self.assertEqual(len(inplace), 0)
        self.assertEqual(inplace.release_after_shape_tagged, [])

    def test_compute_event_action_enum_values(self):
        self.assertEqual(int(si.ComputeEventAction.kInPlace), 0)
        self.assertEqual(int(si.ComputeEventAction.kRelease), 1)
        self.assertEqual(int(si.ComputeEventAction.kReleaseShapeTag), 2)

    def test_compute_context_decision_events(self):
        nodes = [
            oh.make_node("Abs", ["X"], ["A"]),
            oh.make_node("Shape", ["X"], ["S"]),
            oh.make_node("Reshape", ["A", "S"], ["Y"]),
        ]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [2, 3])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [2, 3])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ShapesContext()
        si.compute_shape_model(ctx, model)

        inplace = si.ComputeContext()
        self.assertFalse(inplace.events_enabled)
        self.assertEqual(inplace.events(), [])
        inplace.compute_inplace_reuse_graph(model.graph, ctx, value_tags={"S": "shape"})
        self.assertEqual(inplace.events(), [])

        inplace.events_enabled = True
        inplace.clear_events()
        inplace.compute_inplace_reuse_graph(model.graph, ctx, value_tags={"S": "shape"})

        events = inplace.events()
        self.assertTrue(events)
        self.assertIn(si.ComputeEventAction.kInPlace, [ev.action for ev in events])
        self.assertIn(si.ComputeEventAction.kRelease, [ev.action for ev in events])
        self.assertIn(si.ComputeEventAction.kReleaseShapeTag, [ev.action for ev in events])

        inplace_ev = [
            ev
            for ev in events
            if ev.action == si.ComputeEventAction.kInPlace and ev.node_index == 2
        ]
        self.assertEqual(len(inplace_ev), 1)
        self.assertEqual(inplace_ev[0].output_index, 0)
        self.assertEqual(inplace_ev[0].input_index, 0)
        self.assertEqual(inplace_ev[0].kind, si.InPlaceReuseKind.kEqual)

        release_names = [
            ev.name
            for ev in events
            if ev.action == si.ComputeEventAction.kRelease and ev.node_index == 2
        ]
        self.assertEqual(sorted(release_names), ["A", "S"])

        shape_release = [
            ev
            for ev in events
            if ev.action == si.ComputeEventAction.kReleaseShapeTag and ev.node_index == 2
        ]
        self.assertEqual(len(shape_release), 1)
        self.assertEqual(shape_release[0].name, "S")
        self.assertEqual(shape_release[0].as_dict()["action"], "release_shape_tag")
        self.assertIn("ComputeEvent(action='release_shape_tag'", repr(shape_release[0]))

        inplace.clear_events()
        self.assertEqual(inplace.events(), [])

    def test_compute_orchestration_graph(self):
        nodes = [
            oh.make_node("Abs", ["X"], ["A"]),
            oh.make_node("Abs", ["A"], ["B"]),
            oh.make_node("Abs", ["B"], ["Y"]),
        ]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [3, 4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [3, 4])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ComputeContext()
        ctx.compute(model)

        # Shape inference ran and is held alive by the context.
        self.assertTrue(ctx.shapes.has("A"))
        self.assertTrue(ctx.shapes.has("Y"))
        # In-place reuse was computed for every node.
        self.assertEqual(len(ctx), 3)
        self.assertEqual(self._reuse_pairs(ctx.reuse), [[], [(0, 0)], [(0, 0)]])
        # Peak memory has one entry per node (Abs has no scratch).
        self.assertEqual(list(ctx.peak_memory), [0, 0, 0])
        self.assertEqual(ctx.node_peak_memory(0), 0)

    def test_compute_write_to_model(self):
        nodes = [oh.make_node("Abs", ["X"], ["A"]), oh.make_node("Abs", ["A"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [3, 4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, None)
        model = self._build_model(nodes, [x], [y])

        ctx = si.ComputeContext()
        ctx.compute(model)
        ctx.write_to_model(model)

        # Inferred shape for the intermediate "A" is written into value_info.
        self.assertIn("A", [vi.name for vi in model.graph.value_info])
        # The release-after / in-place metadata is present on the freeing node.
        keys = {p.key for p in model.graph.node[1].metadata_props}
        self.assertIn(si.RELEASE_AFTER_METADATA_KEY, keys)
        self.assertIn(si.INPLACE_REUSE_METADATA_KEY, keys)

    def test_compute_shapes_and_peak_memory_separately(self):
        nodes = [oh.make_node("Abs", ["X"], ["Y"])]
        x = oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [3, 4])
        y = oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, [3, 4])
        model = self._build_model(nodes, [x], [y])

        ctx = si.ComputeContext()
        shapes = ctx.compute_shapes(model)
        self.assertTrue(shapes.has("Y"))
        peak = ctx.compute_peak_memory(model.graph)
        self.assertEqual(list(peak), [0])

    def test_compute_shapes_node_list_with_input_shapes(self):
        nodes = [oh.make_node("Abs", ["X"], ["A"]), oh.make_node("Neg", ["A"], ["Y"])]
        ctx = si.ComputeContext()
        shapes = ctx.compute_shapes(
            nodes, input_shapes={"X": si.SymTensor(onnxl.TensorProto.FLOAT, [3, 4])}
        )
        self.assertTrue(shapes.has("A"))
        self.assertTrue(shapes.has("Y"))
        self.assertEqual(list(shapes.get("Y").shape.dims()), [3, 4])

    def test_compute_shapes_node_list_missing_input_raises(self):
        nodes = [oh.make_node("Abs", ["X"], ["Y"])]
        ctx = si.ComputeContext()
        with self.assertRaises(ValueError):
            ctx.compute_shapes(nodes)

    def test_compute_shapes_function_with_input_shapes(self):
        function = oh.make_function(
            "com.acme",
            "MyFunc",
            ["X"],
            ["Y"],
            [oh.make_node("Abs", ["X"], ["Y"])],
            [oh.make_opsetid("", 18)],
        )
        ctx = si.ComputeContext()
        shapes = ctx.compute_shapes(
            function, input_shapes={"X": si.SymTensor(onnxl.TensorProto.FLOAT, [2, 5])}
        )
        self.assertTrue(shapes.has("Y"))
        self.assertEqual(list(shapes.get("Y").shape.dims()), [2, 5])

    def test_compute_shapes_function_missing_input_raises(self):
        function = oh.make_function(
            "com.acme",
            "MyFunc",
            ["X"],
            ["Y"],
            [oh.make_node("Abs", ["X"], ["Y"])],
            [oh.make_opsetid("", 18)],
        )
        ctx = si.ComputeContext()
        with self.assertRaises(ValueError):
            ctx.compute_shapes(function)


if __name__ == "__main__":
    unittest.main(verbosity=2)
