symbolic_shape_infer.py cannot infer torch.nn.normalize

[Tytskiy]

Describe the issue

BUG

I had tried to infer shape for my model containing torch.nn.functional.normalize, but got error. Even, it happens with x / x.norm(dim=1, keepdim=True)

import torch


class Net(torch.nn.Module):
    def forward(self, x):
        # return torch.nn.functional.normalize(x, dim=1)
        return x / x.norm(dim=1, keepdim=True)

tensor = torch.rand(128, 256)

torch.onnx.export(
    Net(),
    (tensor,),
    'normalize.onnx',
    verbose=True,
    input_names=['tensor',],
    output_names=['output', ],
    dynamic_axes={'tensor': {0: 'batch_size'}},
    opset_version=20,
)

❯ python /home/tytskiy/miniforge3/envs/.../lib/python3.12/site-packages/onnxruntime/tools/symbolic_shape_infer.py --input normalize.onnx
Traceback (most recent call last):
  File "/home/tytskiy/miniforge3/envs/yrec/lib/python3.12/site-packages/onnxruntime/tools/symbolic_shape_infer.py", line 2996, in <module>
    out_mp = SymbolicShapeInference.infer_shapes(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tytskiy/miniforge3/envs/yrec/lib/python3.12/site-packages/onnxruntime/tools/symbolic_shape_infer.py", line 2928, in infer_shapes
    all_shapes_inferred = symbolic_shape_inference._infer_impl()
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tytskiy/miniforge3/envs/yrec/lib/python3.12/site-packages/onnxruntime/tools/symbolic_shape_infer.py", line 2729, in _infer_impl
    out_rank = len(get_shape_from_type_proto(vi.type))
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: object of type 'NoneType' has no len()

normalize.onnx.zip

To reproduce

Just run symbolic_shape_infer.py with onnx model.

Versions:
onnx                      1.17.0                   pypi_0    pypi
onnxruntime-gpu           1.20.1                   pypi_0    pypi
onnxscript                0.1.0.dev20241203          pypi_0    pypi
pytorch                   2.5.1           py3.12_cuda12.4_cudnn9.1.0_0    pytorch

Also, I have tried it with current symbolic_shape_infer.py in repo, but got same

Urgency

No response

Platform

Linux

OS Version

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

1.20.1

ONNX Runtime API

Python

Architecture

X64

Execution Provider

Other / Unknown

Execution Provider Library Version

No response