GroupQueryAttention fails to load when total_sequence_length is computed from static shapes (ParseData type mismatch) #33865

"""ONNX Runtime 1.29: a `GroupQueryAttention` whose `total_sequence_length` is computed from static shapes
(`Squeeze(Cast(Shape(query)[1:2], int32))`) fails to load: shape inference reads it as int64. The same graph with the
value as an int32 constant loads and runs."""

import numpy as np
import onnxruntime as ort
from onnx import TensorProto, helper, numpy_helper

BATCH, LENGTH, HEADS, KV_HEADS, HEAD_SIZE = 1, 7, 4, 2, 16


def build(constant_total: bool) -> bytes:
    nodes = [
        helper.make_node("Shape", ["query"], ["length64"], start=1, end=2),
        helper.make_node("Cast", ["length64"], ["length"], to=TensorProto.INT32),
        helper.make_node("Sub", ["length", "one"], ["last"]),
        helper.make_node("Expand", ["last", "batch"], ["seqlens_k"]),
    ]
    initializers = [
        numpy_helper.from_array(np.array([1], np.int32), "one"),
        numpy_helper.from_array(np.array([BATCH], np.int64), "batch"),
    ]
    if constant_total:
        initializers.append(numpy_helper.from_array(np.array(LENGTH, np.int32), "total"))
    else:
        nodes.append(helper.make_node("Squeeze", ["length"], ["total"]))
    nodes.append(
        helper.make_node(
            "GroupQueryAttention",
            ["query", "key", "value", "", "", "seqlens_k", "total"],
            ["output", "present_key", "present_value"],
            domain="com.microsoft",
            num_heads=HEADS,
            kv_num_heads=KV_HEADS,
        )
    )
    inputs = [
        helper.make_tensor_value_info("query", TensorProto.FLOAT, [BATCH, LENGTH, HEADS * HEAD_SIZE]),
        helper.make_tensor_value_info("key", TensorProto.FLOAT, [BATCH, LENGTH, KV_HEADS * HEAD_SIZE]),
        helper.make_tensor_value_info("value", TensorProto.FLOAT, [BATCH, LENGTH, KV_HEADS * HEAD_SIZE]),
    ]
    outputs = [helper.make_tensor_value_info(name, TensorProto.FLOAT, None) for name in ("output", "present_key", "present_value")]
    graph = helper.make_graph(nodes, "gqa", inputs, outputs, initializers)
    opsets = [helper.make_opsetid("", 20), helper.make_opsetid("com.microsoft", 1)]
    return helper.make_model(graph, opset_imports=opsets, ir_version=10).SerializeToString()


rng = np.random.default_rng(0)
feeds = {
    "query": rng.standard_normal((BATCH, LENGTH, HEADS * HEAD_SIZE)).astype(np.float32),
    "key": rng.standard_normal((BATCH, LENGTH, KV_HEADS * HEAD_SIZE)).astype(np.float32),
    "value": rng.standard_normal((BATCH, LENGTH, KV_HEADS * HEAD_SIZE)).astype(np.float32),
}
print("onnxruntime", ort.__version__)
for constant_total in (False, True):
    label = "total_sequence_length as an int32 constant" if constant_total else "total_sequence_length from Cast(Shape)"
    try:
        session = ort.InferenceSession(build(constant_total), providers=["CPUExecutionProvider"])
        session.run(None, feeds)
        print(f"  {label}: loads and runs")
    except Exception as error:
        print(f"  {label}: {str(error).splitlines()[0][:200]}")