From 95e727be1848e3fa21979210cb49711c20ab5431 Mon Sep 17 00:00:00 2001
From: Andrew White <andrewh@cdw.com>
Date: Fri, 24 Jul 2026 10:56:59 -0500
Subject: [PATCH] fix: allow None dynamic_axes in OnnxLLMExporter.export

OnnxLLMExporter._export_to_onnx unpacked dynamic_axes_input and dynamic_axes_output
directly with {**x, **y}, which raised TypeError when either argument
was left as the default None. Fall back to empty dicts so callers can
omit dynamic axes without crashing.

Adds a unit test that mocks torch.onnx.export and verifies the call
succeeds with dynamic_axes_input=None and dynamic_axes_output=None.

Signed-off-by: Andrew White <andrewh@cdw.com>
---
 nemo_export/onnx_llm_exporter.py              |  2 +-
 .../export/test_onnx_llm_exporter.py          | 20 +++++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/nemo_export/onnx_llm_exporter.py b/nemo_export/onnx_llm_exporter.py
index 60e9ffbe2..94aab556a 100755
--- a/nemo_export/onnx_llm_exporter.py
+++ b/nemo_export/onnx_llm_exporter.py
@@ -236,7 +236,7 @@ def _export_to_onnx(
                 f=self.onnx_model_path,
                 input_names=input_names,
                 output_names=output_names,
-                dynamic_axes={**dynamic_axes_input, **dynamic_axes_output},
+                dynamic_axes={**(dynamic_axes_input or {}), **(dynamic_axes_output or {})},
                 verbose=verbose,
                 opset_version=opset,
                 dynamo=False,
diff --git a/tests/unit_tests/export/test_onnx_llm_exporter.py b/tests/unit_tests/export/test_onnx_llm_exporter.py
index 4075fc0ee..e87b32544 100644
--- a/tests/unit_tests/export/test_onnx_llm_exporter.py
+++ b/tests/unit_tests/export/test_onnx_llm_exporter.py
@@ -106,3 +106,23 @@ def test_quantize_without_modelopt(self):
             pytest.raises(UnavailableError),
         ):
             OnnxLLMExporter().quantize(quant_cfg="", forward_loop="")
+
+    def test_export_allows_none_dynamic_axes(self, temp_dir, dummy_model, dummy_tokenizer):
+        exporter = OnnxLLMExporter(
+            onnx_model_dir=temp_dir,
+            model=dummy_model,
+            tokenizer=dummy_tokenizer,
+            load_runtime=False,
+        )
+        example_inputs = {"input_ids": torch.randint(0, 10, (2, 4), device=exporter.device)}
+        with mock.patch("nemo_export.onnx_llm_exporter.torch.onnx.export") as mock_export:
+            exporter.export(
+                input_names=["input_ids"],
+                output_names=["output"],
+                example_inputs=example_inputs,
+                dynamic_axes_input=None,
+                dynamic_axes_output=None,
+            )
+
+        _, kwargs = mock_export.call_args
+        assert kwargs["dynamic_axes"] == {}