From 1fc58432b6061fb99611a0c04f561d522505c3e8 Mon Sep 17 00:00:00 2001
From: gcunhase <4861122+gcunhase@users.noreply.github.com>
Date: Thu, 5 Feb 2026 14:50:30 -0500
Subject: [PATCH 1/4] Fix: issue with input shape comparison with unknown
 dimensions

Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com>
---
 modelopt/onnx/autocast/referencerunner.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/modelopt/onnx/autocast/referencerunner.py b/modelopt/onnx/autocast/referencerunner.py
index 4eb485a158..220bbd491a 100644
--- a/modelopt/onnx/autocast/referencerunner.py
+++ b/modelopt/onnx/autocast/referencerunner.py
@@ -89,7 +89,12 @@ def _validate_inputs(self, data_loader):
             if sorted(self.input_names) != sorted(data_loader[0].keys()):
                 raise ValueError("Input names from ONNX model do not match provided input names.")
             for inp_name, inp_shape in data_loader[0].items():
-                if self.input_shapes[inp_name] != list(inp_shape.shape):
+                # Get model and data shapes as numpy arrays
+                inp_shape_model = np.array(self.input_shapes[inp_name])
+                inp_shape_data = np.array(inp_shape.shape)
+                # Skip check for unknown dimensions (shape = -1)
+                mask = (inp_shape_model != -1) & (inp_shape_data != -1)
+                if np.any(inp_shape_model[mask] != inp_shape_data[mask]):
                     raise ValueError(
                         f"Input shape from '{inp_name}' does not match provided input shape: "
                         f"{self.input_shapes[inp_name]} vs {list(inp_shape.shape)}. "

From 115fc4810603d10b57fe37f2a27c8962cde044c1 Mon Sep 17 00:00:00 2001
From: gcunhase <4861122+gcunhase@users.noreply.github.com>
Date: Thu, 5 Feb 2026 18:10:12 -0500
Subject: [PATCH 2/4] Add check for input rank

Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com>
---
 modelopt/onnx/autocast/referencerunner.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/modelopt/onnx/autocast/referencerunner.py b/modelopt/onnx/autocast/referencerunner.py
index 220bbd491a..cae4deb086 100644
--- a/modelopt/onnx/autocast/referencerunner.py
+++ b/modelopt/onnx/autocast/referencerunner.py
@@ -92,9 +92,13 @@ def _validate_inputs(self, data_loader):
                 # Get model and data shapes as numpy arrays
                 inp_shape_model = np.array(self.input_shapes[inp_name])
                 inp_shape_data = np.array(inp_shape.shape)
-                # Skip check for unknown dimensions (shape = -1)
-                mask = (inp_shape_model != -1) & (inp_shape_data != -1)
-                if np.any(inp_shape_model[mask] != inp_shape_data[mask]):
+                # Compare input rank
+                raise_value_error = len(inp_shape_model) != len(inp_shape_data)
+                if not raise_value_error:
+                    # Compare input shape, skipping check for unknown dimensions (shape = -1)
+                    mask = (inp_shape_model != -1) & (inp_shape_data != -1)
+                    raise_value_error = np.any(inp_shape_model[mask] != inp_shape_data[mask])
+                if raise_value_error:
                     raise ValueError(
                         f"Input shape from '{inp_name}' does not match provided input shape: "
                         f"{self.input_shapes[inp_name]} vs {list(inp_shape.shape)}. "

From 629e4bbd29f7df6957ff60d7ba893c784db72803 Mon Sep 17 00:00:00 2001
From: gcunhase <4861122+gcunhase@users.noreply.github.com>
Date: Thu, 5 Feb 2026 18:14:14 -0500
Subject: [PATCH 3/4] Add UNK input shape for model shape only

Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com>
---
 modelopt/onnx/autocast/referencerunner.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/modelopt/onnx/autocast/referencerunner.py b/modelopt/onnx/autocast/referencerunner.py
index cae4deb086..18f9cb604e 100644
--- a/modelopt/onnx/autocast/referencerunner.py
+++ b/modelopt/onnx/autocast/referencerunner.py
@@ -96,7 +96,7 @@ def _validate_inputs(self, data_loader):
                 raise_value_error = len(inp_shape_model) != len(inp_shape_data)
                 if not raise_value_error:
                     # Compare input shape, skipping check for unknown dimensions (shape = -1)
-                    mask = (inp_shape_model != -1) & (inp_shape_data != -1)
+                    mask = inp_shape_model != -1
                     raise_value_error = np.any(inp_shape_model[mask] != inp_shape_data[mask])
                 if raise_value_error:
                     raise ValueError(

From 8c26e5b457a78737bde0c2b0d5f67916068944e5 Mon Sep 17 00:00:00 2001
From: gcunhase <4861122+gcunhase@users.noreply.github.com>
Date: Thu, 5 Feb 2026 18:31:45 -0500
Subject: [PATCH 4/4] Update: consider input shape > 0 as static

Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com>
---
 modelopt/onnx/autocast/referencerunner.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/modelopt/onnx/autocast/referencerunner.py b/modelopt/onnx/autocast/referencerunner.py
index 18f9cb604e..27c3bee5d9 100644
--- a/modelopt/onnx/autocast/referencerunner.py
+++ b/modelopt/onnx/autocast/referencerunner.py
@@ -95,8 +95,8 @@ def _validate_inputs(self, data_loader):
                 # Compare input rank
                 raise_value_error = len(inp_shape_model) != len(inp_shape_data)
                 if not raise_value_error:
-                    # Compare input shape, skipping check for unknown dimensions (shape = -1)
-                    mask = inp_shape_model != -1
+                    # Compare input shape, skipping check for unknown dimensions
+                    mask = inp_shape_model > 0
                     raise_value_error = np.any(inp_shape_model[mask] != inp_shape_data[mask])
                 if raise_value_error:
                     raise ValueError(