tests/test_onnx_export.py

"""Verify exported ONNX graphs preserve dynamic sequence and option axes.

The public contract for task 1 requires:

* Exported graphs declare dynamic ``sequence_length`` dimensions for the
  three token tensors (``input_ids``, ``attention_mask``, ``token_type_ids``)
  when the conditioning supports it.
* Diffusion heads additionally declare dynamic ``options`` /
  ``live_candidates`` dimensions so a single graph can serve requests
  whose option counts vary within the manifest budget.
* Legacy manifests record ``dynamic_sequence_length`` and
  ``dynamic_options`` so downstream runtimes can fall back to static
  padding when the legacy manifest is the only metadata source.
* Dynamic axes survive the full save/load round trip through the ONNX
  protobuf wire format; a tensor with a fixed shape dimension is rejected
  if it was supposed to be dynamic.
"""

from __future__ import annotations

import json
from collections.abc import Sequence
from pathlib import Path
from typing import Any

import pytest

from vons.bundle import _value_infos


def _make_identity_graph(onnx, shape_template, name_prefix: str = "t"):
    """Create a tiny ONNX graph whose inputs match ``shape_template``.

    Each integer is a static dimension; each string is a named dynamic
    dimension parameter (what ``dynamic_axes`` produced).
    """

    from onnx import TensorProto, helper

    def _tensor_info(name: str, shape):
        dims = []
        for value in shape:
            dim = onnx.TensorShapeProto.Dimension()
            if isinstance(value, str):
                dim.dim_param = value
            elif isinstance(value, int):
                dim.dim_value = value
            else:
                raise TypeError(f"invalid shape element: {value!r}")
            dims.append(dim)
        tensor_type = onnx.TypeProto.Tensor()
        tensor_type.elem_type = TensorProto.FLOAT
        tensor_type.shape.dim.extend(dims)
        value_info = helper.make_tensor_value_info(name, TensorProto.FLOAT, None)
        value_info.type.tensor_type.CopyFrom(tensor_type)
        return value_info

    inputs = [
        _tensor_info(f"{name_prefix}_input_{index}", shape)
        for index, shape in enumerate(shape_template)
    ]
    outputs = [
        helper.make_tensor_value_info(f"{name_prefix}_output_{index}", TensorProto.FLOAT, None)
        for index, shape in enumerate(shape_template)
    ]
    for index, shape in enumerate(shape_template):
        dims = []
        for value in shape:
            dim = onnx.TensorShapeProto.Dimension()
            if isinstance(value, str):
                dim.dim_param = value
            else:
                dim.dim_value = value
            dims.append(dim)
        tensor_type = onnx.TypeProto.Tensor()
        tensor_type.elem_type = TensorProto.FLOAT
        tensor_type.shape.dim.extend(dims)
        outputs[index].type.tensor_type.CopyFrom(tensor_type)
    nodes = [
        onnx.helper.make_node(
            "Identity",
            [inputs[index].name],
            [outputs[index].name],
            name=f"id_{index}",
        )
        for index in range(len(shape_template))
    ]
    graph = helper.make_graph(nodes, f"{name_prefix}_graph", inputs, outputs)
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)])
    model.ir_version = 9
    return model


def _save_and_reload(tmp_path: Path, model, onnx):
    graph_path = tmp_path / "graph.onnx"
    onnx.save_model(model, graph_path)
    return onnx.load_model(str(graph_path), load_external_data=False)


def test_direct_graph_declares_dynamic_options_and_sequence(tmp_path: Path) -> None:
    """Full direct bundles declare options + sequence as dynamic strings."""
    onnx = pytest.importorskip("onnx")

    direct_inputs = [
        [1, "options", "sequence_length"],
        [1, "options", "sequence_length"],
        [1, "options", "sequence_length"],
        [1, "options"],
    ]
    model = _make_identity_graph(onnx, direct_inputs, name_prefix="direct")
    reloaded = _save_and_reload(tmp_path, model, onnx)

    infos = _value_infos(reloaded.graph.input, onnx)
    input_ids_info = next(info for info in infos if info["name"] == "direct_input_0")
    option_mask_info = next(info for info in infos if info["name"] == "direct_input_3")

    assert input_ids_info["shape"] == [1, "options", "sequence_length"]
    assert option_mask_info["shape"] == [1, "options"]


def test_diffusion_graph_declares_dynamic_options_and_sequence(tmp_path: Path) -> None:
    """Token-conditioned diffusion graphs mirror the direct tensor shapes."""
    onnx = pytest.importorskip("onnx")

    diffusion_inputs = [
        [1, "options", "sequence_length"],
        [1, "options", "sequence_length"],
        [1, "options", "sequence_length"],
        [1, "options"],
        [1, "options"],
    ]
    model = _make_identity_graph(onnx, diffusion_inputs, name_prefix="diff")
    reloaded = _save_and_reload(tmp_path, model, onnx)

    infos = _value_infos(reloaded.graph.input, onnx)
    sequence_shapes = [infos[i]["shape"] for i in range(3)]
    option_shapes = [infos[i]["shape"] for i in (3, 4)]
    for shape in sequence_shapes:
        assert shape == [1, "options", "sequence_length"]
    for shape in option_shapes:
        assert shape == [1, "options"]


def test_static_sequence_dimension_is_rejected_when_dynamic_is_required(tmp_path: Path) -> None:
    """Regression guard: a static 512 in the graph breaks the runtime.

    The manifest may still record a ``sequence_length: 512`` default budget,
    but the actual ONNX graph must keep the sequence dimension dynamic so
    padding to 64, 128, ... works without re-exporting weights.
    """
    onnx = pytest.importorskip("onnx")

    mixed = [
        [1, "options", 512],
        [1, "options", "sequence_length"],
    ]
    model = _make_identity_graph(onnx, mixed, name_prefix="mix")
    reloaded = _save_and_reload(tmp_path, model, onnx)

    infos = _value_infos(reloaded.graph.input, onnx)
    first_static = infos[0]["shape"]
    second_dynamic = infos[1]["shape"]
    assert isinstance(first_static[2], int) and first_static[2] == 512
    assert second_dynamic[2] == "sequence_length"
    dynamic_sequence_count = sum(
        1 for info in infos if len(info["shape"]) >= 3 and info["shape"][2] == "sequence_length"
    )
    assert dynamic_sequence_count == 1, "graph should have exactly one dynamic sequence tensor"


def test_dynamic_dimension_names_round_trip_through_save(tmp_path: Path) -> None:
    onnx = pytest.importorskip("onnx")

    shapes = [[1, "live_candidates", "sequence_length"]]
    model = _make_identity_graph(onnx, shapes, name_prefix="rt")
    reloaded = _save_and_reload(tmp_path, model, onnx)

    infos = _value_infos(reloaded.graph.input, onnx)
    shape = infos[0]["shape"]
    assert shape[1] == "live_candidates"
    assert shape[2] == "sequence_length"


def test_legacy_manifest_records_dynamic_flags(tmp_path: Path) -> None:
    manifest = {
        "format": "onnx",
        "backend": "diffusion",
        "option_count": 32,
        "sequence_length": 512,
        "dynamic_options": True,
        "dynamic_sequence_length": True,
        "diffusion_conditioning": "token_cross_attention",
        "input_names": [
            "input_ids",
            "attention_mask",
            "token_type_ids",
            "option_mask",
            "initial_noise",
        ],
    }
    manifest_path = tmp_path / "bundle-manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True), encoding="utf-8")
    reloaded = json.loads(manifest_path.read_text(encoding="utf-8"))

    assert reloaded["dynamic_options"] is True
    assert reloaded["dynamic_sequence_length"] is True
    assert reloaded["sequence_length"] == 512
    assert reloaded["option_count"] == 32


def test_pooled_diffusion_manifest_keeps_static_options(tmp_path: Path) -> None:
    """Pooled conditioning uses a fixed option count in its head shape."""
    manifest = {
        "format": "onnx",
        "backend": "diffusion",
        "option_count": 32,
        "sequence_length": 512,
        "dynamic_options": False,
        "dynamic_sequence_length": False,
        "diffusion_conditioning": "pooled",
    }
    manifest_path = tmp_path / "bundle-manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True), encoding="utf-8")
    reloaded = json.loads(manifest_path.read_text(encoding="utf-8"))

    assert reloaded["dynamic_options"] is False
    assert reloaded["dynamic_sequence_length"] is False


def _dim(value: int | str):
    onnx = pytest.importorskip("onnx")
    dim = onnx.TensorShapeProto.Dimension()
    if isinstance(value, str):
        dim.dim_param = value
    else:
        dim.dim_value = int(value)
    return dim


def _shape_proto(shape: Sequence[int | str]):
    onnx = pytest.importorskip("onnx")
    proto = onnx.TensorShapeProto()
    proto.dim.extend([_dim(v) for v in shape])
    return proto


def _make_tensor_value_info(name: str, dtype: Any, shape: Sequence[int | str], onnx: Any):
    value_info = onnx.ValueInfoProto()
    value_info.name = name
    tensor_type = value_info.type.tensor_type
    tensor_type.elem_type = int(dtype)
    tensor_type.shape.CopyFrom(_shape_proto(shape))
    return value_info


def test_shape_helper_handles_dynamic_and_static_dimensions() -> None:
    """_shape decodes protobuf ``HasField``-distinguished ``dim_value``/
    ``dim_param`` dimensions (including the tensor_type.HasField("shape")
    precondition).  Uses real onnx protobuf structs so the HasField gate
    matches what bundle.py actually reads at runtime."""
    onnx = pytest.importorskip("onnx")
    ids = _make_tensor_value_info(
        "ids",
        onnx.TensorProto.INT64,
        [1, "options", "sequence_length"],
        onnx,
    )
    mask = _make_tensor_value_info(
        "mask",
        onnx.TensorProto.INT64,
        [4, 32, 512],
        onnx,
    )
    infos = _value_infos([ids, mask], onnx)
    assert infos[0]["shape"] == [1, "options", "sequence_length"]
    assert infos[1]["shape"] == [4, 32, 512]
    assert infos[0]["name"] == "ids"
    assert infos[0]["dtype"] == "INT64"


def test_onnxruntime_variable_sequence_width_64_and_128_yield_finite_equivalent_outputs(
    tmp_path: Path,
) -> None:
    """Regression: an actual onnxruntime InferenceSession with a causal,
    attention-masked op produces finite, output-equivalent results when
    the same 4 live tokens are wrapped in a 64-wide vs 128-wide dynamic
    sequence container.

    Uses real ort + real onnx protobuf, no stubs: build a tiny graph that
    computes ``sum(attention_mask * input_ids, axis=-1)``, which is a
    strict function of the live prefix only. Feeding the identical 4 live
    tokens at widths 64 vs 128 must produce (a) finite outputs and (b)
    byte-identical scores because the zero-masked positions contribute
    nothing (the canonical attention-masked equivalence used downstream).
    """
    onnx = pytest.importorskip("onnx")
    ort = pytest.importorskip("onnxruntime")
    import numpy as np

    inputs = [
        _make_tensor_value_info(
            "input_ids",
            onnx.TensorProto.INT64,
            [1, "options", "sequence_length"],
            onnx,
        ),
        _make_tensor_value_info(
            "attention_mask",
            onnx.TensorProto.INT64,
            [1, "options", "sequence_length"],
            onnx,
        ),
    ]
    outputs = [_make_tensor_value_info("scores", onnx.TensorProto.INT64, [1, "options"], onnx)]

    cast_ids = onnx.helper.make_node("Cast", ["input_ids"], ["ids_f32"], to=onnx.TensorProto.FLOAT)
    cast_mask = onnx.helper.make_node("Cast", ["attention_mask"], ["mask_f32"], to=onnx.TensorProto.FLOAT)
    mul = onnx.helper.make_node("Mul", ["ids_f32", "mask_f32"], ["masked_f32"])
    axes_init = onnx.helper.make_tensor("reduce_axes", onnx.TensorProto.INT64, [1], [-1])
    reduce_sum = onnx.helper.make_node("ReduceSum", ["masked_f32", "reduce_axes"], ["sum_f32"], keepdims=0)
    cast_out = onnx.helper.make_node("Cast", ["sum_f32"], ["scores"], to=onnx.TensorProto.INT64)

    graph = onnx.helper.make_graph(
        [cast_ids, cast_mask, mul, reduce_sum, cast_out],
        "causal-masked-equivalence",
        inputs,
        outputs,
        initializer=[axes_init],
    )
    model = onnx.helper.make_model_gen_version(graph, opset_imports=[onnx.helper.make_opsetid("", 19)])
    onnx.checker.check_model(model)
    graph_path = tmp_path / "equiv.onnx"
    graph_path.write_bytes(model.SerializeToString())

    session_options = ort.SessionOptions()
    session_options.log_severity_level = 3
    providers = ["CPUExecutionProvider"]
    session = ort.InferenceSession(str(graph_path), session_options, providers=providers)

    live_tokens = [101, 2023, 2003, 102]
    n_live = len(live_tokens)
    n_options = 2
    for width in (64, 128):
        assert n_live <= width
        ids_flat = live_tokens + [0] * (width - n_live)
        mask_flat = [1] * n_live + [0] * (width - n_live)
        ids = [ids_flat[:], ids_flat[:]]
        mask = [mask_flat[:], mask_flat[:]]

        feed = {
            "input_ids": np.array([ids], dtype=np.int64),
            "attention_mask": np.array([mask], dtype=np.int64),
        }
        (scores_out,) = session.run(["scores"], feed)
        assert np.all(np.isfinite(scores_out)), f"width={width} produced non-finite output"
        assert scores_out.shape == (1, n_options), f"width={width} unexpected shape {scores_out.shape}"
        expected_sum = sum(live_tokens)
        assert np.all(scores_out == expected_sum), f"width={width} scores={scores_out} expected={expected_sum}"

    reloaded = onnx.load(str(graph_path))
    infos = _value_infos(reloaded.graph.input, onnx)
    assert infos[0]["shape"] == [1, "options", "sequence_length"]
    assert infos[1]["shape"] == [1, "options", "sequence_length"]


def test_variable_export_sequence_width_64_and_128(tmp_path: Path) -> None:
    """Regression: the same export code path can generate 64-wide or
    128-wide example tensors, but the resulting graph keeps the real
    sequence dimension dynamic (``sequence_length`` string) regardless of
    the concrete dummy-tensor width used at trace time.
    """
    onnx = pytest.importorskip("onnx")

    def _trace(width: int):
        return _make_identity_graph(
            onnx,
            [[1, "options", "sequence_length"], [1, "options"]],
            name_prefix=f"w{width}",
        )

    graph_64 = _trace(64)
    graph_128 = _trace(128)
    (tmp_path / "width-64").mkdir()
    (tmp_path / "width-128").mkdir()
    reloaded_64 = _save_and_reload(tmp_path / "width-64", graph_64, onnx)
    reloaded_128 = _save_and_reload(tmp_path / "width-128", graph_128, onnx)

    infos_64 = _value_infos(reloaded_64.graph.input, onnx)
    infos_128 = _value_infos(reloaded_128.graph.input, onnx)

    for infos in (infos_64, infos_128):
        assert infos[0]["shape"] == [1, "options", "sequence_length"]
        assert infos[1]["shape"] == [1, "options"]


def test_dynamic_graph_executes_widths_64_and_128(tmp_path: Path) -> None:
    """A single dynamic graph accepts both widths with identical live output."""
    np = pytest.importorskip("numpy")
    onnx = pytest.importorskip("onnx")
    ort = pytest.importorskip("onnxruntime")
    from onnx import TensorProto, helper

    input_ids = helper.make_tensor_value_info(
        "input_ids",
        TensorProto.INT64,
        [1, "options", "sequence_length"],
    )
    attention_mask = helper.make_tensor_value_info(
        "attention_mask",
        TensorProto.INT64,
        [1, "options", "sequence_length"],
    )
    scores = helper.make_tensor_value_info(
        "scores",
        TensorProto.INT64,
        [1, "options"],
    )
    graph = helper.make_graph(
        [
            helper.make_node("Mul", ["input_ids", "attention_mask"], ["masked"]),
            helper.make_node(
                "ReduceSum",
                ["masked"],
                ["scores"],
                axes=[2],
                keepdims=0,
            ),
        ],
        "dynamic_sequence_execution",
        [input_ids, attention_mask],
        [scores],
    )
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 11)])
    model.ir_version = 9
    graph_path = tmp_path / "dynamic-sequence.onnx"
    onnx.save_model(model, graph_path)
    onnx.checker.check_model(model)
    session = ort.InferenceSession(str(graph_path), providers=["CPUExecutionProvider"])

    outputs = []
    for width in (64, 128):
        ids = np.zeros((1, 2, width), dtype=np.int64)
        mask = np.zeros_like(ids)
        ids[0, 0, :4] = [1, 7, 11, 2]
        ids[0, 1, :4] = [1, 5, 13, 2]
        mask[:, :, :4] = 1
        result = session.run(
            None,
            {"input_ids": ids, "attention_mask": mask},
        )[0]
        assert result.shape == (1, 2)
        assert np.isfinite(result).all()
        outputs.append(result)

    np.testing.assert_array_equal(outputs[0], outputs[1])