builder_test.py

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from __future__ import annotations

import ast
import unittest
from typing import Sequence

import onnx_ir as ir

import onnxscript._internal.builder as builder
import onnxscript.testing
from onnxscript import script
from onnxscript.onnx_types import DOUBLE, FLOAT, INT64

_default_opset_version = 23


def _resolve_type_spec(spec: builder.TypeSpec) -> ir.TypeAndShape:
    """Convert a *TypeSpec* to an :class:`ir.TypeAndShape`.

    Accepts either an :class:`ir.TypeAndShape` directly, or a
    :class:`~onnxscript.onnx_types.TensorType` subclass (e.g. ``FLOAT[1024]``
    or ``FLOAT['M', 'N']``).

    NOTE: This is a local copy of :func:`builder._resolve_type_spec` so that
    tests do not reference a private helper directly.
    """
    from onnxscript.onnx_types import TensorType  # pylint: disable=import-outside-toplevel

    if isinstance(spec, ir.TypeAndShape):
        return spec
    if isinstance(spec, type) and issubclass(spec, TensorType):
        return spec.to_ir_type_and_shape()
    raise TypeError(f"Expected ir.TypeAndShape or a TensorType subclass, got {type(spec)!r}.")


def _build(
    input_types: Sequence[builder.TypeSpec],
    trace_function=None,
    output_types: Sequence[builder.TypeSpec] | None = None,
) -> ir.Graph:
    graph = ir.Graph(
        name="test_model",
        inputs=[],
        outputs=[],
        nodes=[],
        opset_imports={"": _default_opset_version},
    )

    resolved_inputs = [_resolve_type_spec(t) for t in input_types]
    for i, ts in enumerate(resolved_inputs):
        graph.inputs.append(ir.Value(name=f"input_{i}", type=ts.type, shape=ts.shape))

    if trace_function is not None:
        graph_builder = builder.GraphBuilder(graph)
        outputs = trace_function(graph_builder.op, *graph.inputs)
        if not isinstance(outputs, Sequence):
            outputs = [outputs]

        if output_types is not None:
            resolved_outputs = [_resolve_type_spec(t) for t in output_types]
            if len(outputs) != len(resolved_outputs):
                raise ValueError(
                    f"Expected {len(resolved_outputs)} outputs, but got {len(outputs)}."
                )
            for output, ts in zip(outputs, resolved_outputs):
                output.type = ts.type
                output.merge_shapes(ts.shape)

        graph.outputs.extend(outputs)

    return graph


def _create_builder_with_inputs() -> tuple[builder.OpBuilder, ir.Value, ir.Value]:
    """Create a graph builder with two float tensor inputs (shape [2, 3, 4]).

    Returns:
        A tuple of (op_builder, input_x, input_y).
    """
    graph = _build(input_types=[FLOAT[2, 3, 4], FLOAT[2, 3, 4]])
    graph_builder = builder.GraphBuilder(graph)
    x, y = graph.inputs
    return graph_builder.op, x, y


class GraphBuilderTest(unittest.TestCase):
    def test_builder_basic(self):
        def _add_mul_add(op: builder.OpBuilder, x: ir.Value, y: ir.Value) -> ir.Value:
            t1 = op.Add(x, y)
            t2 = op.Mul(x, y)
            z = op.Add(t1, t2)
            return z

        float_2d = ir.TypeAndShape(ir.TensorType(ir.DataType.FLOAT), ir.Shape([3, 4]))
        graph = _build(
            input_types=[float_2d, float_2d],
            trace_function=_add_mul_add,
            output_types=[float_2d],
        )
        # Expect exactly 3 nodes: Add, Mul, Add
        op_types = [node.op_type for node in graph]
        self.assertEqual(op_types, ["Add", "Mul", "Add"])

        # Verify inputs and outputs
        self.assertEqual(len(graph.inputs), 2)
        self.assertEqual(len(graph.outputs), 1)

        # Verify the connectivity: final Add takes outputs of the first Add and Mul
        nodes = list(graph)
        add1, mul, add2 = nodes
        self.assertEqual(list(add2.inputs), [add1.outputs[0], mul.outputs[0]])

    def test_value_naming(self):
        """Test that output names can be specified via the _outputs option."""

        def _add_with_custom_names(
            op: builder.OpBuilder, x: ir.Value, y: ir.Value
        ) -> ir.Value:
            # Specify custom names for output values
            t1 = op.Add(x, y, _outputs=["add_result"])
            t2 = op.Mul(x, y, _outputs=["mul_result"])
            z = op.Add(t1, t2, _outputs=["final_add"])
            return z

        float_2d = ir.TypeAndShape(ir.TensorType(ir.DataType.FLOAT), ir.Shape([3, 4]))
        graph = _build(
            input_types=[float_2d, float_2d],
            trace_function=_add_with_custom_names,
            output_types=[float_2d],
        )

        # Verify that the nodes have outputs with the specified names
        nodes = list(graph)
        self.assertEqual(len(nodes), 3)

        # Check output names (v_ prefix is added to all value names)
        self.assertEqual(nodes[0].outputs[0].name, "v_add_result")
        self.assertEqual(nodes[1].outputs[0].name, "v_mul_result")
        self.assertEqual(nodes[2].outputs[0].name, "v_final_add")

        # Verify the final output has the correct name
        self.assertEqual(len(graph.outputs), 1)
        self.assertEqual(graph.outputs[0].name, "v_final_add")

    def test_value_naming_with_hierarchy(self):
        """Test that hierarchical naming works with user-specified output names."""
        op, x, y = _create_builder_with_inputs()

        # Test custom names at root level
        t1 = op.Add(x, y, _outputs=["my_add"])
        self.assertEqual(t1.name, "v_my_add")

        # Test custom names with hierarchical context
        op.builder.push_module("layer1")
        t2 = op.Mul(t1, y, _outputs=["my_mul"])
        self.assertEqual(t2.name, "v_layer1.my_mul")

        # Test nested hierarchical context with custom names
        op.builder.push_module("attention")
        t3 = op.Add(t2, x, _outputs=["my_nested_add"])
        self.assertEqual(t3.name, "v_layer1.attention.my_nested_add")

        # Pop back and verify prefix is applied correctly
        op.builder.pop_module()
        t4 = op.Mul(t3, y, _outputs=["another_mul"])
        self.assertEqual(t4.name, "v_layer1.another_mul")

        op.builder.pop_module()
        t5 = op.Add(t4, x, _outputs=["final_result"])
        self.assertEqual(t5.name, "v_final_result")

    def test_value_naming_with_ir_value_objects(self):
        """Test that hierarchical naming works when passing ir.Value objects as _outputs."""
        op, x, y = _create_builder_with_inputs()

        # Create pre-named ir.Value objects
        out1 = ir.Value(name="my_output")
        out2 = ir.Value(name="layer_output")
        out3 = ir.Value(name="nested_output")

        # Test at root level
        t1 = op.Add(x, y, _outputs=[out1])
        self.assertEqual(t1.name, "v_my_output")
        self.assertIs(t1, out1)

        # Test with hierarchical context
        op.builder.push_module("layer1")
        t2 = op.Mul(t1, y, _outputs=[out2])
        self.assertEqual(t2.name, "v_layer1.layer_output")
        self.assertIs(t2, out2)

        # Test nested hierarchical context
        op.builder.push_module("attention")
        t3 = op.Add(t2, x, _outputs=[out3])
        self.assertEqual(t3.name, "v_layer1.attention.nested_output")
        self.assertIs(t3, out3)

    def test_default_output_naming_strategy(self):
        """Test the default naming strategy for generated output values using op_type_nX_output format."""

        def _ops_with_default_names(
            op: builder.OpBuilder, x: ir.Value, y: ir.Value
        ) -> ir.Value:
            # Single output operations should be named {op_type}_nX_output where X is node count
            t1 = op.Add(x, y)
            t2 = op.Mul(x, y)
            z = op.Add(t1, t2)
            return z

        float_2d = ir.TypeAndShape(ir.TensorType(ir.DataType.FLOAT), ir.Shape([3, 4]))
        graph = _build(
            input_types=[float_2d, float_2d],
            trace_function=_ops_with_default_names,
            output_types=[float_2d],
        )

        # Verify the nodes use the new naming strategy
        nodes = list(graph)
        self.assertEqual(len(nodes), 3)

        # Check output names follow the v_{op_type}_{count} pattern for single outputs
        self.assertEqual(nodes[0].outputs[0].name, "v_Add_0")
        self.assertEqual(nodes[1].outputs[0].name, "v_Mul_1")
        self.assertEqual(nodes[2].outputs[0].name, "v_Add_2")

        # Verify the final output has the correct name
        self.assertEqual(len(graph.outputs), 1)
        self.assertEqual(graph.outputs[0].name, "v_Add_2")

    def test_hierarchical_naming(self):
        """Test the hierarchical naming strategy (for value and node names)."""
        op, x, y = _create_builder_with_inputs()

        # Test node and value naming at root level
        t1 = op.Add(x, y)
        self.assertEqual(t1.name, "v_Add_0")
        self.assertEqual(t1.producer().name, "Add_node_0")

        t2 = op.Mul(t1, y)
        self.assertEqual(t2.name, "v_Mul_1")
        self.assertEqual(t2.producer().name, "Mul_node_1")

        # Test node and value naming with hierarchical context prefix
        op.builder.push_module("layer1")
        t3 = op.Add(t2, x)
        self.assertEqual(t3.name, "v_layer1.Add_2")
        self.assertEqual(t3.producer().name, "layer1/Add_node_2")

        # Test nested hierarchical context
        op.builder.push_module("attention")
        t4 = op.Mul(t3, y)
        self.assertEqual(t4.name, "v_layer1.attention.Mul_3")
        self.assertEqual(t4.producer().name, "layer1/attention/Mul_node_3")

        # Pop back to layer1 and verify naming continues correctly
        op.builder.pop_module()
        t5 = op.Add(t4, x)
        self.assertEqual(t5.name, "v_layer1.Add_4")
        self.assertEqual(t5.producer().name, "layer1/Add_node_4")

        # Pop back to root context
        op.builder.pop_module()
        t6 = op.Mul(t5, y)
        self.assertEqual(t6.name, "v_Mul_5")
        self.assertEqual(t6.producer().name, "Mul_node_5")

    def test_shape_inference_add(self):
        """Test that shape inference works correctly for Add operation."""
        op, x, y = _create_builder_with_inputs()

        # Create Add node without explicitly setting output type/shape
        result = op.Add(x, y)

        # Verify output type is inferred correctly
        self.assertIsNotNone(result.type)
        self.assertEqual(result.type.dtype, ir.DataType.FLOAT)

        # Verify output shape is inferred correctly
        self.assertIsNotNone(result.shape)
        self.assertEqual(list(result.shape), [2, 3, 4])

        self.assertEqual(result.name, "v_Add_0")

    def test_custom_domain_explicit(self):
        """Test using operations from custom domains with explicit _domain parameter."""
        op, x, y = _create_builder_with_inputs()

        # Create a custom domain operation with explicit _domain parameter
        # Using "com.microsoft" as an example domain
        result = op.CustomOp(x, y, _domain="com.microsoft")

        # Verify the node was created with the correct domain
        nodes = list(op.builder.graph)
        self.assertEqual(len(nodes), 1)
        node = nodes[0]
        self.assertEqual(node.domain, "com.microsoft")
        self.assertEqual(node.op_type, "CustomOp")

        # Verify inputs and outputs are connected correctly
        self.assertEqual(list(node.inputs), [x, y])
        self.assertEqual(node.outputs[0], result)

    def test_custom_domain_with_version(self):
        """Test using operations from custom domains with explicit _domain and _version parameters."""
        op, x, y = _create_builder_with_inputs()

        # Create a custom domain operation with explicit _domain and _version parameters
        result = op.MicrosoftOp(x, y, _domain="com.microsoft", _version=10)

        # Verify the node was created with the correct domain and version
        nodes = list(op.builder.graph)
        self.assertEqual(len(nodes), 1)
        node = nodes[0]
        self.assertEqual(node.domain, "com.microsoft")
        self.assertEqual(node.op_type, "MicrosoftOp")
        self.assertEqual(node.version, 10)

        # Verify output value is created
        self.assertIsNotNone(result)
        self.assertEqual(result.name, "v_MicrosoftOp_0")

    def test_multiple_custom_domain_operations(self):
        """Test mixing operations from multiple domains."""
        op, x, y = _create_builder_with_inputs()

        # Create standard domain operation
        t1 = op.Add(x, y)

        # Create custom domain operation
        t2 = op.CustomOp(t1, y, _domain="com.microsoft")

        # Create another custom domain operation with different domain
        _ = op.AnotherOp(t2, x, _domain="com.custom")

        # Verify all nodes were created with correct domains
        nodes = list(op.builder.graph)
        self.assertEqual(len(nodes), 3)

        self.assertEqual(nodes[0].domain, "")
        self.assertEqual(nodes[0].op_type, "Add")

        self.assertEqual(nodes[1].domain, "com.microsoft")
        self.assertEqual(nodes[1].op_type, "CustomOp")

        self.assertEqual(nodes[2].domain, "com.custom")
        self.assertEqual(nodes[2].op_type, "AnotherOp")

    def test_opset_builder_for_custom_domain(self):
        """Test creating and using an opset builder for a custom domain."""
        op, x, y = _create_builder_with_inputs()

        # Create an OpBuilder for the "com.microsoft" domain with version 1
        ms_op = op.builder.opset("com.microsoft", 1)

        # Use operations through the custom domain opset builder
        t1 = ms_op.CustomOp(x, y)
        _ = ms_op.AnotherOp(t1, x)

        # Verify all nodes were created with the correct domain
        nodes = list(op.builder.graph)
        self.assertEqual(len(nodes), 2)

        # Verify first operation
        self.assertEqual(nodes[0].domain, "com.microsoft")
        self.assertEqual(nodes[0].op_type, "CustomOp")
        self.assertEqual(nodes[0].version, 1)
        self.assertEqual(list(nodes[0].inputs), [x, y])

        # Verify second operation
        self.assertEqual(nodes[1].domain, "com.microsoft")
        self.assertEqual(nodes[1].op_type, "AnotherOp")
        self.assertEqual(nodes[1].version, 1)
        self.assertEqual(list(nodes[1].inputs), [t1, x])

    def test_mixed_domain_opsets(self):
        """Test using both standard domain and custom domain opset builders together."""
        op, x, y = _create_builder_with_inputs()

        # Create custom domain opset builder
        ms_op = op.builder.opset("com.microsoft", 2)

        # Mix operations from different domains
        t1 = op.Add(x, y)  # Standard domain operation
        t2 = ms_op.MsAdd(t1, y)  # Custom domain operation
        _ = op.Mul(t2, x)  # Back to standard domain

        # Verify nodes were created with correct domains
        nodes = list(op.builder.graph)
        self.assertEqual(len(nodes), 3)

        self.assertEqual(nodes[0].domain, "")
        self.assertEqual(nodes[0].op_type, "Add")

        self.assertEqual(nodes[1].domain, "com.microsoft")
        self.assertEqual(nodes[1].op_type, "MsAdd")
        self.assertEqual(nodes[1].version, 2)

        self.assertEqual(nodes[2].domain, "")
        self.assertEqual(nodes[2].op_type, "Mul")

    def test_scalar_constant_initializers_are_cached(self):
        """Test that scalar constants produce named initializers and are shared across nodes."""
        graph = ir.Graph(
            name="test_model",
            inputs=[],
            outputs=[],
            nodes=[],
            opset_imports={"": _default_opset_version},
        )
        # Create one int64 input and one float32 input
        x = ir.Value(name="x", type=ir.TensorType(ir.DataType.INT64), shape=ir.Shape([3]))
        y = ir.Value(name="y", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([3]))
        graph.inputs.extend([x, y])

        graph_builder = builder.GraphBuilder(graph)
        op = graph_builder.op

        # Two Adds that both use the integer constant 1
        r1 = op.Add(x, 1)
        r2 = op.Add(x, 1)
        # Two Adds that both use the float constant 1.0
        r3 = op.Add(y, 1.0)
        r4 = op.Add(y, 1.0)

        # The two int Add nodes should share the same constant initializer (same ir.Value)
        int_const_1 = r1.producer().inputs[1]
        int_const_2 = r2.producer().inputs[1]
        self.assertIs(int_const_1, int_const_2)
        self.assertEqual(int_const_1.name, "const_1_i64")

        # The two float Add nodes should share the same constant initializer
        float_const_1 = r3.producer().inputs[1]
        float_const_2 = r4.producer().inputs[1]
        self.assertIs(float_const_1, float_const_2)
        self.assertEqual(float_const_1.name, "const_1.0_f32")

        # The int and float constants should be different ir.Values
        self.assertIsNot(int_const_1, float_const_1)

    def test_int_constant_cast_to_float_via_like_type(self):
        """Test that op.Add(float_x, 1) converts the int 1 to a float tensor.

        When the schema binds the int constant to the same type variable as a float
        input, the constant should be created with dtype FLOAT and named accordingly.
        """
        graph = ir.Graph(
            name="test_model",
            inputs=[],
            outputs=[],
            nodes=[],
            opset_imports={"": _default_opset_version},
        )
        x = ir.Value(name="x", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([3]))
        graph.inputs.append(x)

        graph_builder = builder.GraphBuilder(graph)
        op = graph_builder.op

        _ = op.Add(x, 1)

        nodes = list(graph)
        self.assertEqual(len(nodes), 1)

        # The int constant 1 should have been cast to float and named with f32 suffix
        const_input = nodes[0].inputs[1]
        self.assertEqual(const_input.name, "const_1_f32")
        # The constant's type should be FLOAT, not INT64
        self.assertEqual(const_input.const_value.dtype, ir.DataType.FLOAT)

    def test_int_constant_with_unknown_type_uses_cast_like(self):
        """Test that op.Add(unknown_x, 1) produces an int tensor + CastLike.

        When the input type is unknown, the constant is created with its natural
        Python type (int -> i64), and a CastLike node is inserted to dynamically
        cast it at runtime.
        """
        graph = ir.Graph(
            name="test_model",
            inputs=[],
            outputs=[],
            nodes=[],
            opset_imports={"": _default_opset_version},
        )
        # Input with no type information
        x = ir.Value(name="x", shape=ir.Shape([3]))
        graph.inputs.append(x)

        graph_builder = builder.GraphBuilder(graph)
        op = graph_builder.op

        _ = op.Add(x, 1)

        nodes = list(graph)
        # Expect 2 nodes: CastLike (to cast the int constant to x's type) + Add
        self.assertEqual(len(nodes), 2)

        cast_like_node = nodes[0]
        add_node = nodes[1]

        self.assertEqual(cast_like_node.op_type, "CastLike")
        self.assertEqual(add_node.op_type, "Add")

        # The original constant should be int64-typed
        const_initializer = cast_like_node.inputs[0]
        self.assertEqual(const_initializer.name, "const_1_i64")
        self.assertEqual(const_initializer.const_value.dtype, ir.DataType.INT64)

        # CastLike's second input should be x (the like_type reference)
        self.assertIs(cast_like_node.inputs[1], x)

        # Add should use the CastLike output, not the raw constant
        self.assertIs(add_node.inputs[1], cast_like_node.outputs[0])

    def test_pop_module_raises_on_empty_stack(self):
        """Test that pop_module raises RuntimeError when no module has been pushed."""
        op, _, _ = _create_builder_with_inputs()

        # Popping without any push should raise
        with self.assertRaises(RuntimeError):
            op.builder.pop_module()

        # Push then pop is fine; a second pop should raise
        op.builder.push_module("layer1")
        op.builder.pop_module()
        with self.assertRaises(RuntimeError):
            op.builder.pop_module()

    def test_output_names_are_unique_for_same_op_type(self):
        """Test that repeated calls to the same op produce unique output names."""
        op, x, y = _create_builder_with_inputs()

        t1 = op.Add(x, y)
        t2 = op.Add(x, y)
        t3 = op.Add(x, y)

        # Each Add output should have a unique name via the node count suffix
        self.assertEqual(t1.name, "v_Add_0")
        self.assertEqual(t2.name, "v_Add_1")
        self.assertEqual(t3.name, "v_Add_2")

        # Verify all names are distinct
        names = [t1.name, t2.name, t3.name]
        self.assertEqual(len(set(names)), 3)

    def test_multi_output_names_are_unique(self):
        """Test that multi-output ops produce unique names with counter suffix."""
        op, x, y = _create_builder_with_inputs()

        # First multi-output call
        out1_a, out1_b = op.TopK(x, 1, axis=-1, _outputs=2)
        # Second multi-output call
        out2_a, out2_b = op.TopK(y, 1, axis=-1, _outputs=2)

        # Each call should produce unique names
        self.assertNotEqual(out1_a.name, out2_a.name)
        self.assertNotEqual(out1_b.name, out2_b.name)

    def test_node_metadata_props_namespace(self):
        """Test that nodes have namespace metadata matching the scope hierarchy."""
        op, x, y = _create_builder_with_inputs()

        # Root-level node
        t1 = op.Add(x, y)
        self.assertEqual(t1.producer().metadata_props["namespace"], "")

        # Node inside a module scope
        op.builder.push_module("layer1", "DecoderLayer")
        t2 = op.Mul(t1, y)
        self.assertEqual(t2.producer().metadata_props["namespace"], "layer1: DecoderLayer")

        # Nested scope
        op.builder.push_module("self_attn", "Attention")
        t3 = op.Add(t2, x)
        self.assertEqual(
            t3.producer().metadata_props["namespace"],
            "layer1: DecoderLayer/self_attn: Attention",
        )
        op.builder.pop_module()
        op.builder.pop_module()

    def test_node_metadata_props_class_hierarchy(self):
        """Test that nodes have class hierarchy metadata."""
        op, x, y = _create_builder_with_inputs()

        op.builder.push_module("layer1", "DecoderLayer")
        op.builder.push_module("self_attn", "Attention")
        t1 = op.MatMul(x, y)
        node = t1.producer()

        self.assertEqual(
            node.metadata_props["pkg.onnxscript.class_hierarchy"],
            repr(["DecoderLayer", "Attention"]),
        )
        self.assertEqual(
            node.metadata_props["pkg.onnxscript.name_scopes"],
            repr(["layer1", "self_attn"]),
        )
        # class_hierarchy and name_scopes have the same length
        self.assertEqual(
            len(ast.literal_eval(node.metadata_props["pkg.onnxscript.class_hierarchy"])),
            len(ast.literal_eval(node.metadata_props["pkg.onnxscript.name_scopes"])),
        )
        op.builder.pop_module()
        op.builder.pop_module()

    def test_attributes_are_created_properly(self):
        """Test that int, float, str, and list attributes are set correctly on a node."""
        op, x, y = _create_builder_with_inputs()

        result = op.DummyOp(
            x,
            y,
            _domain="test.domain",
            int_attr=42,
            float_attr=3.14,
            str_attr="hello",
            ints_attr=[1, 2, 3],
            floats_attr=[1.0, 2.0, 3.0],
            strs_attr=["a", "b", "c"],
        )

        node = result.producer()
        self.assertEqual(node.op_type, "DummyOp")
        self.assertEqual(node.domain, "test.domain")

        # Verify scalar attributes
        int_attr = node.attributes["int_attr"]
        self.assertEqual(int_attr.type, ir.AttributeType.INT)
        self.assertEqual(int_attr.value, 42)

        float_attr = node.attributes["float_attr"]
        self.assertEqual(float_attr.type, ir.AttributeType.FLOAT)
        self.assertAlmostEqual(float_attr.value, 3.14)

        str_attr = node.attributes["str_attr"]
        self.assertEqual(str_attr.type, ir.AttributeType.STRING)
        self.assertEqual(str_attr.value, "hello")

        # Verify list attributes
        ints_attr = node.attributes["ints_attr"]
        self.assertEqual(ints_attr.type, ir.AttributeType.INTS)
        self.assertEqual(list(ints_attr.value), [1, 2, 3])

        floats_attr = node.attributes["floats_attr"]
        self.assertEqual(floats_attr.type, ir.AttributeType.FLOATS)
        self.assertEqual(list(floats_attr.value), [1.0, 2.0, 3.0])

        strs_attr = node.attributes["strs_attr"]
        self.assertEqual(strs_attr.type, ir.AttributeType.STRINGS)
        self.assertEqual(list(strs_attr.value), ["a", "b", "c"])

    def test_call_inlines_onnxscript_function(self):
        """Test that GraphBuilder.call inlines an @onnxscript.script function."""
        # Create a GraphBuilder first
        op, x, y = _create_builder_with_inputs()

        # Define the script function after creating op, using op as default_opset
        @script(default_opset=op)
        def mul_add_relu(X, Y):
            tmp = X * Y
            tmp = tmp + X
            return op.Relu(tmp)

        result = op.call(mul_add_relu, x, y)

        # The inlined function should produce 3 nodes: Mul, Add, Relu
        nodes = list(op.builder.graph)
        op_types = [n.op_type for n in nodes]
        self.assertEqual(op_types, ["Mul", "Add", "Relu"])

        # The result should be a single ir.Value (the Relu output)
        self.assertIsInstance(result, ir.Value)

        # Verify connectivity: Relu takes the Add output
        relu_node = nodes[2]
        add_node = nodes[1]
        self.assertIs(relu_node.inputs[0], add_node.outputs[0])

        # Verify the Add takes the Mul output and original input x
        mul_node = nodes[0]
        self.assertIs(add_node.inputs[0], mul_node.outputs[0])
        self.assertIs(add_node.inputs[1], x)

        # Verify the Mul takes the original inputs x and y
        self.assertIs(mul_node.inputs[0], x)
        self.assertIs(mul_node.inputs[1], y)

    def test_call_with_outputs_option(self):
        """Test that GraphBuilder.call respects the _outputs option for renaming."""
        # Create a GraphBuilder first
        op, x, y = _create_builder_with_inputs()

        # Define the script function after creating op, using op as default_opset
        @script(default_opset=op)
        def add_mul(X, Y):
            a = X + Y
            b = X * Y
            return a, b

        result = op.call(add_mul, x, y, _outputs=["sum_result", "product_result"])

        # The result should be a list of 2 ir.Values (when function returns multiple outputs)
        self.assertIsInstance(result, list)
        self.assertEqual(len(result), 2)
        sum_result, product_result = result

        # Verify output names are correctly set (with v_ prefix from value naming convention)
        self.assertEqual(sum_result.name, "v_sum_result")
        self.assertEqual(product_result.name, "v_product_result")

        # Verify the nodes were created correctly
        nodes = list(op.builder.graph)
        self.assertEqual(len(nodes), 2)
        self.assertEqual(nodes[0].op_type, "Add")
        self.assertEqual(nodes[1].op_type, "Mul")

    def test_call_with_outer_scope_value(self):
        """Test that script supports references to pre-existing values."""
        # Create a GraphBuilder first
        op, x, y = _create_builder_with_inputs()
        product = op.Mul(x, y)

        @script()
        def add_product(X):
            return op.Add(X, product)  # Reference to 'product' from outer scope

        x_plus = op.call(add_product, x, _outputs=["x_plus"])
        y_plus = op.call(add_product, y, _outputs=["y_plus"])

        op.builder.graph.outputs.extend([x_plus, y_plus])

        # Now, create the same graph directly:
        op2, x2, y2 = _create_builder_with_inputs()
        product2 = op2.Mul(x2, y2)
        x2_plus = op2.Add(x2, product2, _outputs=["x_plus"])
        y2_plus = op2.Add(y2, product2, _outputs=["y_plus"])
        op2.builder.graph.outputs.extend([x2_plus, y2_plus])

        # Verify that the two graphs are structurally equivalent
        onnxscript.testing.assert_isomorphic_graph(op.builder.graph, op2.builder.graph)

    def test_call_with_prefix_option(self):
        """Test that GraphBuilder.call respects the _prefix option for hierarchical naming."""
        # Create a GraphBuilder first
        op, x, y = _create_builder_with_inputs()

        # Define the script function after creating op, using op as default_opset
        @script(default_opset=op)
        def mul_add_relu(X, Y):
            tmp = X * Y
            tmp = tmp + X
            return op.Relu(tmp)

        result = op.call(mul_add_relu, x, y, _prefix="layer1")

        # The nodes should have the prefix in their names
        nodes = list(op.builder.graph)
        self.assertEqual(len(nodes), 3)

        # Check that all node names start with the prefix (node names use / separator)
        for node in nodes:
            self.assertTrue(
                node.name.startswith("layer1/"),
                f"Node name {node.name} should start with layer1/",
            )

        # Verify the result is a single ir.Value
        self.assertIsInstance(result, ir.Value)

    def test_call_with_outputs_and_prefix_options(self):
        """Test that GraphBuilder.call respects both _outputs and _prefix options together.

        Note: _outputs names are set before the prefix context is applied, so they don't get
        the prefix in their names. However, the inlined nodes do get the prefix applied, and
        intermediate values (not renamed by _outputs) do get the prefix applied.
        """
        # Create a GraphBuilder first
        op, x, y = _create_builder_with_inputs()

        # Define the script function after creating op, using op as default_opset
        @script(default_opset=op)
        def add_mul(X, Y):
            # Intermediate values that are not explicitly renamed by _outputs
            XSquare = X * X
            YSquare = Y * Y
            # Final outputs that will be renamed by _outputs
            a = XSquare + Y
            b = XSquare * YSquare
            return a, b

        result = op.call(
            add_mul, x, y, _outputs=["custom_sum", "custom_product"], _prefix="math_ops"
        )

        # The result should be a list of 2 ir.Values
        self.assertIsInstance(result, list)
        self.assertEqual(len(result), 2)
        sum_result, product_result = result

        # Verify output names are set (with v_ prefix from value naming convention)
        self.assertEqual(sum_result.name, "v_custom_sum")
        self.assertEqual(product_result.name, "v_custom_product")

        # Verify all nodes have the prefix applied to their names
        nodes = list(op.builder.graph)
        self.assertEqual(len(nodes), 4)  # Mul (XSquare), Mul (YSquare), Add, Mul (final)

        # All node names should start with prefix (node names use / separator)
        for node in nodes:
            self.assertTrue(
                node.name.startswith("math_ops/"),
                f"Node name {node.name} should start with math_ops/",
            )

        # Verify intermediate value names also get the prefix (value names use v_ prefix and . separator)
        # The first Mul produces XSquare
        x_square = nodes[0].outputs[0]
        self.assertTrue(
            x_square.name.startswith("v_math_ops."),
            f"Intermediate value {x_square.name} should have prefix",
        )

        # The second Mul produces YSquare
        y_square = nodes[1].outputs[0]
        self.assertTrue(
            y_square.name.startswith("v_math_ops."),
            f"Intermediate value {y_square.name} should have prefix",
        )

    def test_call_outputs_mismatch_error(self):
        """Test that GraphBuilder.call raises an error if _outputs has wrong count."""
        # Create a GraphBuilder first
        op, x, y = _create_builder_with_inputs()

        # Define the script function after creating op, using op as default_opset
        @script(default_opset=op)
        def add_mul(X, Y):
            a = X + Y
            b = X * Y
            return a, b

        # The function returns 2 outputs, but we provide only 1 name
        with self.assertRaises(ValueError) as cm:
            op.call(add_mul, x, y, _outputs=["only_one_name"])

        self.assertIn("does not match", str(cm.exception))


class BuildSubgraphTest(unittest.TestCase):
    """Tests for GraphBuilder.subgraph()."""

    def _make_builder(self, opset_version: int = 23) -> builder.GraphBuilder:
        """Return a minimal GraphBuilder for the given opset version."""
        graph = ir.Graph(
            name="parent",
            inputs=[],
            outputs=[],
            nodes=[],
            opset_imports={"": opset_version},
        )
        return builder.GraphBuilder(graph)

    def test_basic_subgraph(self):
        """Subgraph returns a valid ir.Graph with correct inputs/outputs."""

        def _add(op, x, y):
            return op.Add(x, y)

        gb = self._make_builder()
        graph = gb.subgraph(
            _add,
            inputs=[FLOAT[3, 4], FLOAT[3, 4]],
            outputs=[FLOAT[3, 4]],
        )
        self.assertIsInstance(graph, ir.Graph)
        self.assertEqual(len(graph.inputs), 2)
        self.assertEqual(len(graph.outputs), 1)
        op_types = [node.op_type for node in graph]
        self.assertEqual(op_types, ["Add"])

    def test_subgraph_inherits_opset_version(self):
        """The subgraph opset version matches the parent GraphBuilder."""
        gb = self._make_builder(opset_version=17)
        graph = gb.subgraph(
            lambda op, x: op.Identity(x),
            inputs=[FLOAT[...]],
            outputs=[FLOAT[...]],
        )
        self.assertEqual(graph.opset_imports[""], 17)

    def test_subgraph_with_ir_type_and_shape(self):
        """Subgraph also accepts ir.TypeAndShape directly."""

        def _mul(op, x, y):
            return op.Mul(x, y)

        float_2d = ir.TypeAndShape(ir.TensorType(ir.DataType.FLOAT), ir.Shape([2, 3]))
        gb = self._make_builder()
        graph = gb.subgraph(
            _mul,
            inputs=[float_2d, float_2d],
            outputs=[float_2d],
        )
        self.assertIsInstance(graph, ir.Graph)
        self.assertEqual(len(list(graph)), 1)
        self.assertEqual(next(iter(graph)).op_type, "Mul")

    def test_subgraph_multiple_outputs(self):
        """Subgraph handles multiple outputs."""

        def _add_and_mul(op, x, y):
            return op.Add(x, y), op.Mul(x, y)

        ts = FLOAT[...]
        gb = self._make_builder()
        graph = gb.subgraph(
            _add_and_mul,
            inputs=[ts, ts],
            outputs=[ts, ts],
        )
        self.assertEqual(len(graph.outputs), 2)

    def test_subgraph_output_count_mismatch_raises(self):
        """Subgraph raises ValueError when output count does not match."""

        def _returns_one(op, x, y):
            return op.Add(x, y)

        gb = self._make_builder()
        with self.assertRaises(ValueError):
            gb.subgraph(
                _returns_one,
                inputs=[FLOAT[...], FLOAT[...]],
                outputs=[FLOAT[...], FLOAT[...]],  # expects 2, gets 1
            )

    def test_subgraph_custom_name(self):
        """Subgraph passes the name through to the ir.Graph."""

        def _id(op, x):
            return op.Identity(x)

        gb = self._make_builder()
        graph = gb.subgraph(
            _id,
            inputs=[DOUBLE[...]],
            outputs=[DOUBLE[...]],
            name="scan_body",
        )
        self.assertEqual(graph.name, "scan_body")

    def test_invalid_type_spec_raises(self):
        """Subgraph raises TypeError for an unrecognised type specification."""

        def _id(op, x):
            return op.Identity(x)

        gb = self._make_builder()
        with self.assertRaises(TypeError):
            gb.subgraph(
                _id,
                inputs=["not_a_type_spec"],
                outputs=["not_a_type_spec"],
            )

    def test_subgraph_dict_inputs_outputs(self):
        """Subgraph accepts a dict to name inputs and outputs."""

        def _add(op, x, y):
            return op.Add(x, y)

        gb = self._make_builder()
        graph = gb.subgraph(
            _add,
            inputs={"x": FLOAT[3, 4], "y": FLOAT[3, 4]},
            outputs={"sum": FLOAT[3, 4]},
        )
        self.assertIsInstance(graph, ir.Graph)
        self.assertEqual(len(graph.inputs), 2)
        self.assertEqual(graph.inputs[0].name, "x")
        self.assertEqual(graph.inputs[1].name, "y")
        self.assertEqual(len(graph.outputs), 1)
        self.assertEqual(graph.outputs[0].name, "sum")

    def test_subgraph_list_auto_names(self):
        """List-based inputs/outputs get auto-generated names."""

        def _id(op, x):
            return op.Identity(x)

        gb = self._make_builder()
        graph = gb.subgraph(
            _id,
            inputs=[FLOAT[...]],
            outputs=[FLOAT[...]],
        )
        self.assertEqual(graph.inputs[0].name, "input_0")
        self.assertEqual(graph.outputs[0].name, "output_0")


class BuildGraphFunctionTest(unittest.TestCase):
    """Tests for the module-level build_graph() utility."""

    def test_build_graph_basic(self):
        """build_graph works without a parent GraphBuilder."""
        graph = builder.build_graph(
            lambda op, x, y: op.Add(x, y),
            inputs={"x": FLOAT[3, 4], "y": FLOAT[3, 4]},
            outputs={"sum": FLOAT[3, 4]},
            opset_imports={"": 20},
        )
        self.assertIsInstance(graph, ir.Graph)
        self.assertEqual(graph.opset_imports[""], 20)
        self.assertEqual(graph.inputs[0].name, "x")
        self.assertEqual(graph.inputs[1].name, "y")
        self.assertEqual(graph.outputs[0].name, "sum")

    def test_build_graph_custom_name(self):
        """build_graph passes name to the ir.Graph."""
        graph = builder.build_graph(
            lambda op, x: op.Identity(x),
            inputs=[FLOAT[...]],
            outputs=[FLOAT[...]],
            name="loop_body",
        )
        self.assertEqual(graph.name, "loop_body")


class PartitionInputsAttributesTest(unittest.TestCase):
    """Tests for GraphBuilder._partition_inputs_attributes."""

    def test_unknown_op_passes_inputs_and_kwargs_through(self):
        """An unknown op has no schema, so inputs and kwargs pass through unchanged."""

        def _dummy(op, x, y):
            return op.DummyOp(x, y, alpha=1.0)

        graph = _build(
            input_types=[FLOAT[3, 4], FLOAT[3, 4]],
            trace_function=_dummy,
        )
        x, y = graph.inputs
        node = graph.node(0)
        self.assertEqual(node.op_type, "DummyOp")
        self.assertEqual(list(node.inputs), [x, y])
        self.assertEqual(node.attributes["alpha"].as_float(), 1.0)

    def test_op_with_only_inputs(self):
        """Add has two inputs and no attributes."""

        def _add(op, x, y):
            return op.Add(x, y)

        graph = _build(
            input_types=[FLOAT[3, 4], FLOAT[3, 4]],
            trace_function=_add,
        )
        x, y = graph.inputs
        node = graph.node(0)
        self.assertEqual(node.op_type, "Add")
        self.assertEqual(list(node.inputs), [x, y])
        self.assertEqual(len(node.attributes), 0)

    def test_op_with_inputs_and_attributes_in_kwargs(self):
        """Gemm has 3 inputs (A, B, C) and attributes (alpha, beta, transA, transB)."""

        def _gemm(op, a, b, c):
            return op.Gemm(a, b, c, alpha=2.0, transB=1)

        graph = _build(
            input_types=[FLOAT[3, 4], FLOAT[4, 5], FLOAT[3, 5]],
            trace_function=_gemm,
        )
        a, b, c = graph.inputs
        node = graph.node(0)
        self.assertEqual(node.op_type, "Gemm")
        self.assertEqual(list(node.inputs), [a, b, c])
        self.assertEqual(node.attributes["alpha"].as_float(), 2.0)
        self.assertEqual(node.attributes["transB"].as_int(), 1)

    def test_op_with_optional_input_omitted(self):
        """Gemm's third input (C) is optional. Omitting it should work."""

        def _gemm_no_c(op, a, b):
            return op.Gemm(a, b, alpha=2.0)

        graph = _build(
            input_types=[FLOAT[3, 4], FLOAT[4, 5]],
            trace_function=_gemm_no_c,
        )
        a, b = graph.inputs
        node = graph.node(0)
        self.assertEqual(node.op_type, "Gemm")
        self.assertEqual(list(node.inputs), [a, b])
        self.assertEqual(node.attributes["alpha"].as_float(), 2.0)

    def test_does_not_fill_attribute_defaults(self):
        """Attribute defaults should not be filled in (fill_defaults=False)."""

        def _gemm_no_attrs(op, a, b):
            return op.Gemm(a, b)

        graph = _build(
            input_types=[FLOAT[3, 4], FLOAT[4, 5]],
            trace_function=_gemm_no_attrs,
        )
        node = graph.node(0)
        # alpha, beta, transA, transB all have defaults but should NOT appear
        self.assertFalse(node.attributes)

    def test_variadic_inputs_with_attribute(self):
        """Concat has variadic inputs and an axis attribute."""

        def _concat(op, x, y, z):
            return op.Concat(x, y, z, axis=0)

        graph = _build(
            input_types=[FLOAT[3, 4], FLOAT[3, 4], FLOAT[3, 4]],
            trace_function=_concat,
        )
        x, y, z = graph.inputs
        node = graph.node(0)
        self.assertEqual(node.op_type, "Concat")
        self.assertEqual(list(node.inputs), [x, y, z])
        self.assertEqual(node.attributes["axis"].as_int(), 0)

    def test_slice_kwargs_are_correctly_ordered_as_inputs(self):
        """Calling op.Slice with keyword arguments should place them in schema order."""

        def _slice(op, data, starts, ends, axes, steps):
            # Pass optional inputs as kwargs in non-schema order
            return op.Slice(data, ends=ends, steps=steps, starts=starts, axes=axes)

        graph = _build(
            input_types=[FLOAT[20, 10], INT64[2], INT64[2], INT64[2], INT64[2]],
            trace_function=_slice,
        )
        data, starts, ends, axes, steps = graph.inputs

        slice_node = graph.node(0)
        self.assertEqual(slice_node.op_type, "Slice")
        # Schema order: data, starts, ends, axes, steps
        self.assertEqual(list(slice_node.inputs), [data, starts, ends, axes, steps])

    def test_omitting_required_input_raises(self):
        """Omitting a required input should raise TypeError."""

        def _add_missing_input(op, x):
            return op.Add(x)

        with self.assertRaises(TypeError):
            _build(
                input_types=[FLOAT[3, 4]],
                trace_function=_add_missing_input,
            )

    def test_extra_inputs_raises(self):
        """Extra positional inputs beyond the schema should raise TypeError."""

        def _add_extra_input(op, x, y, z):
            return op.Add(x, y, z)

        with self.assertRaises(TypeError):
            _build(
                input_types=[FLOAT[3, 4], FLOAT[3, 4], FLOAT[3, 4]],
                trace_function=_add_extra_input,
            )


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