Repository: pytorch
Issue: INTERNAL ASSERT FAILED in `shape_type_inference.cpp`
Labels: oncall: jit

### 🐛 Describe the bug

I made a slight modification to a PyTorch model that previously exported to ONNX without issue (specifically, instead of using the passed in input `mesh_edge_indices`, the modified model loads them from a file - I assume this is converted into a constant tensor when tracing).

I first `torch.jit.trace`d and saved the model to produce `postcvpr.pt` (put inside a .zip folder and [attached](https://github.com/pytorch/pytorch/files/12540209/postcvpr.zip)). However, when trying to export this new model to ONNX I hit an `INTERNAL ASSERT FAILED`.

For reference, the unmodified model (which does not trigger the assertion) is attached [here](https://github.com/pytorch/pytorch/files/12541034/postcvpr_old.zip).

Python code used to export to ONNX:
```python
from typing import Sequence
import torch
from torch.jit._trace import TopLevelTracedModule

traced_model: TopLevelTracedModule = torch.jit.load("postcvpr.pt")

INPUT_NAMES = [
    "cloth_features",
    "active_mask",
    "obstacle_features",
    "mesh_edge_indices",
    "mesh_edge_features",
    "coarse0_edge_indices",
    "coarse0_edge_features",
    "coarse1_edge_indices",
    "coarse1_edge_features",
    "coarse2_edge_indices",
    "coarse2_edge_features",
    "inverse_world_edge_indices",
    "inverse_world_edge_features",
    "direct_world_edge_indices",
    "direct_world_edge_features",
]

INPUT_SHAPES = [
    (torch.float32, (4424, 24)),
    (torch.bool, (5678, 1)),
    (torch.float32, (5678, 24)),
    (torch.int64, (2, 26272)),
    (torch.float32, (26272, 12)),
    (torch.int64, (2, 10568)),
    (torch.float32, (10568, 12)),
    (torch.int64, (2, 6340)),
    (torch.float32, (6340, 12)),
    (torch.int64, (2, 2900)),
    (torch.float32, (2900, 12)),
    (torch.int64, (2, 782)),
    (torch.float32, (782, 9)),
    (torch.int64, (2, 782)),
    (torch.float32, (782, 9)),
]

# Adapted from https://github.com/onnx/onnx/issues/654
def export_onnx_model(
    model,
    input_shapes: list[tuple[torch.dtype, Sequence[int]]],
    onnx_path,
    input_names=None,
    output_names=None,
    dynamic_axes=None,
):
    # Remove the old onnx model - make sure we are actually creating a new one, lol
    # os.remove(onnx_path)

    inputs = create_inputs(input_shapes)

    # Test just running the model with generated sample input
    model(*inputs)

    # Export
    torch.onnx.export(
        model,
        inputs,
        onnx_path,
        input_names=input_names,
        output_names=output_names,
        dynamic_axes=dynamic_axes,
        opset_version=16,
    )


def create_inputs(input_shapes: list[tuple[torch.dtype, Sequence[int]]]):
    # Minimum of either number of cloth verts or number of body verts
    min_verts = min(input_shapes[0][1][0], input_shapes[1][1][0])
    return tuple(
        map(
            lambda s: (
                ty := s[0],
                shape := s[1],
                torch.rand(shape, device="cuda") < 0.9
                if ty == torch.bool
                else torch.randint(0, min_verts, shape, dtype=ty, device="cuda")
                if ty == torch.int64
                else torch.rand(shape, dtype=ty, device="cuda"),
            )[-1],
            input_shapes,
        )
    )

export_onnx_model(
    traced_model,
    INPUT_SHAPES,
    "./postcvpr.onnx",
    INPUT_NAMES,
    ["output"],
)
```

Full Output:
```
/home/nathaniel/miniconda3/envs/hood/lib/python3.10/site-packages/torch/onnx/utils.py:825: UserWarning: no signature found for <torch.ScriptMethod object at 0x7fb6d36f8860>, skipping _decide_input_format
  warnings.warn(f"{e}, skipping _decide_input_format")
============= Diagnostic Run torch.onnx.export version 2.0.1+cu117 =============
verbose: False, log level: Level.ERROR
======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================

Traceback (most recent call last):
  File "/home/nathaniel/dev/HOOD_2/HOOD/./Test.py", line 90, in <module>
    export_onnx_model(
  File "/home/nathaniel/dev/HOOD_2/HOOD/./Test.py", line 61, in export_onnx_model
    torch.onnx.export(
  File "/home/nathaniel/miniconda3/envs/hood/lib/python3.10/site-packages/torch/onnx/utils.py", line 506, in export
    _export(
  File "/home/nathaniel/miniconda3/envs/hood/lib/python3.10/site-packages/torch/onnx/utils.py", line 1548, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "/home/nathaniel/miniconda3/envs/hood/lib/python3.10/site-packages/torch/onnx/utils.py", line 1117, in _model_to_graph
    graph = _optimize_graph(
  File "/home/nathaniel/miniconda3/envs/hood/lib/python3.10/site-packages/torch/onnx/utils.py", line 665, in _optimize_graph
    graph = _C._jit_pass_onnx(graph, operator_export_type)
  File "/home/nathaniel/miniconda3/envs/hood/lib/python3.10/site-packages/torch/onnx/utils.py", line 1891, in _run_symbolic_function
    return symbolic_fn(graph_context, *inputs, **attrs)
  File "/home/nathaniel/miniconda3/envs/hood/lib/python3.10/site-packages/torch/onnx/symbolic_helper.py", line 392, in wrapper
    return fn(g, *args, **kwargs)
  File "/home/nathaniel/miniconda3/envs/hood/lib/python3.10/site-packages/torch/onnx/symbolic_opset9.py", line 945, in expand_as
    return g.op("Expand", self, shape)
  File "/home/nathaniel/miniconda3/envs/hood/lib/python3.10/site-packages/torch/onnx/_internal/jit_utils.py", line 86, in op
    return _add_op(self, opname, *raw_args, outputs=outputs, **kwargs)
  File "/home/nathaniel/miniconda3/envs/hood/lib/python3.10/site-packages/torch/onnx/_internal/jit_utils.py", line 245, in _add_op
    node = _create_node(
  File "/home/nathaniel/miniconda3/envs/hood/lib/python3.10/site-packages/torch/onnx/_internal/jit_utils.py", line 306, in _create_node
    _C._jit_pass_onnx_node_shape_type_inference(node, params_dict, opset_version)
RuntimeError: input_shape_value == reshape_value || input_shape_value == 1 || reshape_value == 1 INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/onnx/shape_type_inference.cpp":580, please report a bug to PyTorch. ONNX Expand input shape constraint not satisfied.
```

### Versions

```
Collecting environment information...
PyTorch version: 2.0.1+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 13.1.0-8ubuntu1~22.04) 13.1.0
Clang version: Could not collect
CMake version: version 3.27.1
Libc version: glibc-2.35

Python version: 3.10.9 (main, Mar  8 2023, 10:47:38) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.90.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: Quadro T2000 with Max-Q Design
Nvidia driver version: 531.41
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Address sizes:                   39 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          12
On-line CPU(s) list:             0-11
Vendor ID:                       GenuineIntel
Model name:                      Intel(R) Xeon(R) W-10855M CPU @ 2.80GHz
CPU family:                      6
Model:                           165
Thread(s) per core:              2
Core(s) per socket:              6
Socket(s):                       1
Stepping:                        2
BogoMIPS:                        5615.99
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology cpuid pni pclmulqdq ssse3 fma cx16 pdcm pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt xsaveopt xsavec xgetbv1 xsaves flush_l1d arch_capabilities
Hypervisor vendor:               Microsoft
Virtualization type:             full
L1d cache:                       192 KiB (6 instances)
L1i cache:                       192 KiB (6 instances)
L2 cache:                        1.5 MiB (6 instances)
L3 cache:                        12 MiB (1 instance)
Vulnerability Itlb multihit:     KVM: Mitigation: VMX unsupported
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Retbleed:          Mitigation; Enhanced IBRS
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:             Unknown: Dependent on hypervisor status
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] numpy==1.22.4
[pip3] pytorch3d==0.7.4
[pip3] torch==2.0.1
[pip3] torch-cluster==1.6.1
[pip3] torch-geometric==2.3.1
[pip3] torch-scatter==2.1.1
[pip3] torch-sparse==0.6.17
[pip3] torchaudio==2.0.2
[pip3] torchinfo==1.8.0
[pip3] torchvision==0.15.2
[pip3] triton==2.0.0
[conda] blas                      1.0                         mkl    conda-forge
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2023.1.0         h6d00ec8_46342  
[conda] mkl-service               2.4.0           py310h5eee18b_1  
[conda] mkl_fft                   1.3.6           py310h1128e8f_1  
[conda] mkl_random                1.2.2           py310h1128e8f_1  
[conda] mxnet-mkl                 1.6.0                    pypi_0    pypi
[conda] numpy                     1.22.4                   pypi_0    pypi
[conda] pytorch                   2.0.1           py3.10_cuda11.7_cudnn8.5.0_0    pytorch
[conda] pytorch-cuda              11.7                 h778d358_5    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] pytorch3d                 0.7.4           py310_cu117_pyt201    pytorch3d
[conda] torch                     2.0.1                    pypi_0    pypi
[conda] torch-cluster             1.6.1                    pypi_0    pypi
[conda] torch-geometric           2.3.1                    pypi_0    pypi
[conda] torch-scatter             2.1.1                    pypi_0    pypi
[conda] torch-sparse              0.6.17                   pypi_0    pypi
[conda] torchaudio                2.0.2               py310_cu117    pytorch
[conda] torchinfo                 1.8.0                    pypi_0    pypi
[conda] torchtriton               2.0.0                     py310    pytorch
[conda] torchvision               0.15.2              py310_cu117    pytorch
[conda] triton                    2.0.0                    pypi_0    pypi
```

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel