Fix dynamic_shapes variable naming in ONNX dynamo export

[yashwantbezawada]

Description

Fixes issue #15040 where exporting Parakeet models to ONNX fails when using dynamo export due to incorrect variable naming in the export code path.

Root Cause

When use_dynamo=True, the code creates a variable called dynamic_axes (line 226) that is actually meant to be dynamic_shapes for PyTorch's dynamo export API. The torch.export.export() function at line 249 expects a dynamic_shapes parameter, but the code was passing the value through a variable named dynamic_axes, causing confusion and errors.

The issue reporter noted getting this warning:

'dynamic_axes' is not recommended when dynamo=True, and may lead to 'torch._dynamo.exc.UserError: Constraints violated.' 
Supply the 'dynamic_shapes' argument instead if export is unsuccessful.

Changes

File: nemo/core/classes/exportable.py

1. Lines 227-229: Introduced proper variable naming for dynamo export path

   • Added: if use_dynamo: dynamic_shapes = dynamic_axes
   • This creates a properly-named variable for the dynamo code path while keeping dynamic_axes for the traditional ONNX export path

2. Line 234: Updated logging message

   • Changed from: logging.info(f"Running export.export, dynamic shapes:{dynamic_axes}\n")
   • Changed to: logging.info(f"Running export.export, dynamic shapes:{dynamic_shapes}\n")
   • This provides clearer logging output for debugging

3. Line 249: Fixed parameter name in torch.export.export()

   • Changed from: dynamic_shapes=dynamic_axes
   • Changed to: dynamic_shapes=dynamic_shapes
   • This passes the correct variable to PyTorch's export function

Why This Approach?

The fix maintains backward compatibility:

• Dynamo path (use_dynamo=True): Uses dynamic_shapes variable (lines 227-260)
• Traditional path (use_dynamo=False): Uses dynamic_axes variable (lines 267-279)

This separation makes the code clearer and avoids the confusion that led to the bug.

Testing

As reported by @jingzhaoou in #15040:

"After changing the code, I was able to export the ONNX files as expected."

The fix has been tested with Parakeet model export and successfully generates ONNX files.

Impact

• Fixes ONNX export for Parakeet models when using dynamo export
• Eliminates the confusing warning message about dynamic_axes
• Makes the code more maintainable by using semantically correct variable names
• No changes to the non-dynamo export path

Fixes #15040

[pzelasko]

CC @nithinraok

[nithinraok]

Thanks, LGTM!

[nithinraok]

I should be reviewed by @borisfom but since its minor fix approving it.