Batch-transpose fusion into `FusedMatMul` sets `transBatchA` on inputs of different ranks, which ONNX Runtime can't compute

Abierto
  • ¿Algo no está bien?
#3,156 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 4 días

@Syedmudassir1604 ya está trabajando en esto.

Desde el 5/11/2026.

  • #3159 de @Syedmudassir1604 — abierto

Evaluación

Dificultad
3/5
Tiempo estimado
1-2 días
Aptitud para principiantes
72/100
Tipo de issue
Error
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
python
Área
backend, machine-learning

Línea de trabajo

Start by finding the ORT fusion rules that fold Transpose into FusedMatMul as transBatchA or transBatchB, then trace how they inspect input shapes and ranks. Reproduce the reported case with rank 3 x and rank 2 w, and check that the fusion is skipped when the ranks differ. Done means the optimized graph runs in ONNX Runtime and matches the unfused result.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

module: optimizer

The rules that fold a Transpose into FusedMatMul as transBatchA / transBatchB don't check the other input's rank. ONNX Runtime's FusedMatMul only batch-transposes when both inputs have the same rank, of at least 3, so the fused graph fails at runtime. A projection of a transposed activation, Transpose(x, [1, 0, 2]) @ w with x rank 3 and w rank 2, is enough.

Repro (onnxscript 0.7.2, onnxruntime 1.31.1):

"""onnxscript folds a batch `Transpose` into `FusedMatMul` (`transBatchA`) when the other input has a lower rank, which
ONNX Runtime can't compute.

`Div(MatMul(Transpose(x, perm=[1, 0, 2]), w), 2)` with `x` rank 3 and `w` rank 2 (a projection of a transposed
activation). onnxscript's ORT fusions turn it into `FusedMatMul(x, w, transBatchA=1, alpha=0.5)`.
"""

import numpy as np
import onnx_ir as ir
import onnxruntime as ort
from onnx import TensorProto, helper
from onnxscript.rewriter.ort_fusions import optimize_for_ort

x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [4, 2, 8])
w = helper.make_tensor("w", TensorProto.FLOAT, [8, 16], np.random.default_rng(0).standard_normal((8, 16)).astype(np.float32).ravel())
two = helper.make_tensor("two", TensorProto.FLOAT, [], [2.0])
nodes = [
    helper.make_node("Transpose", ["x"], ["xt"], perm=[1, 0, 2]),
    helper.make_node("MatMul", ["xt", "w"], ["product"]),
    helper.make_node("Div", ["product", "two"], ["y"]),
]
y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 4, 16])
proto = helper.make_model(helper.make_graph(nodes, "g", [x], [y], [w, two]), opset_imports=[helper.make_opsetid("", 20)], ir_version=10)

feeds = {"x": np.random.default_rng(1).standard_normal((4, 2, 8)).astype(np.float32)}
expected = ort.InferenceSession(proto.SerializeToString(), providers=["CPUExecutionProvider"]).run(None, feeds)[0]

model, _ = optimize_for_ort(ir.serde.deserialize_model(proto))
for node in model.graph.all_nodes():
    print(node.op_type, {k: a.value for k, a in node.attributes.items()}, [str(v.shape) for v in node.inputs if v is not None])
try:
    got = ort.InferenceSession(ir.to_proto(model).SerializeToString(), providers=["CPUExecutionProvider"]).run(None, feeds)[0]
    print("fused graph max diff:", np.abs(got - expected).max())
except Exception as error:
    print("ONNX Runtime:", str(error).split("Status Message: ")[-1].strip())

FusedMatMul {'alpha': 0.5, 'transBatchA': 1} ['[4,2,8]', '[8,16]']
ONNX Runtime: .../matmul_helper.h:73 ... left_num_dims > 2 && left_num_dims == right_num_dims was false. Two inputs should have same rank and rank >= 3 if transBatchA or transBatchB is true

The unfused graph runs and gives [2, 4, 16].

Lenguaje dominante
Python
Estrellas
457
Forks
149
Merge medio
3 d 20 h
PR fusionados (30 d)
22

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de microsoft/onnxscript

Todos los issues de microsoft/onnxscript

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.