Repository: pytorch
Issue: Custom symbolic functions for ONNX export with None args causes SEGFAULT
Labels: module: crash,module: onnx,triaged

### 🐛 Describe the bug

`torch.onn.export` with custom symbolic function produces SEGFAULT.

It happens when I set some aguments optional with `None`. OnnxRuntime opset allows it https://github.com/microsoft/onnxruntime/blob/3e4c5e64877c6d9814e4ebce5dcbb1fe71588ec5/docs/ContribOperators.md#commicrosoftpackedmultiheadattention

**Repro:**

```python
import torch

class Function(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x: torch.Tensor, cu_seqlens: torch.Tensor, token_offset: torch.Tensor) -> torch.Tensor:
        if torch.onnx.is_in_onnx_export():
            # doesn't matter – onnx only check number of ouputs and their shapes
            return x

        # do something
        ...

    @staticmethod
    def symbolic(g: torch.Graph, x: torch.Value, cu_seqlens: torch.Value, token_offset: torch.Value) -> torch.Value:
        return g.op(
            'com.microsoft::PackedMultiHeadAttention',
            x,
            None,
            None,
            None,
            token_offset,
            cu_seqlens,
            None,
            num_heads_i=1,
        ).setType(x.type())


class Net(torch.nn.Module):
    def forward(self, x: torch.Tensor, cu_seqlens: torch.Tensor, token_offset: torch.Tensor) -> torch.Tensor:
        return Function.apply(x, cu_seqlens, token_offset)


net = Net()

embeddings = torch.tensor([10])
token_offset = torch.tensor([10])
cu_seqlens = torch.tensor([10])

torch.onnx.export(
    net,
    (embeddings, cu_seqlens, token_offset),
    'graph.onnx',
    input_names=['embeddings', 'cu_seqlens', 'lengths'],
    dynamo=False,
    verbose=True,
    opset_version=20,
)

```

**Output:**

As is:
```bash
❯ python example.py
[1]    4159913 segmentation fault  python example.py
```

with GDB:
```bash
❯ gdb python
GNU gdb (Ubuntu 12.1-0ubuntu1~22.04.2) 12.1
Copyright (C) 2022 Free Software Foundation, Inc.
License GPLv3+: GNU GPL version 3 or later <http://gnu.org/licenses/gpl.html>
This is free software: you are free to change and redistribute it.
There is NO WARRANTY, to the extent permitted by law.
Type "show copying" and "show warranty" for details.
This GDB was configured as "x86_64-linux-gnu".
Type "show configuration" for configuration details.
For bug reporting instructions, please see:
<https://www.gnu.org/software/gdb/bugs/>.
Find the GDB manual and other documentation resources online at:
    <http://www.gnu.org/software/gdb/documentation/>.

For help, type "help".
Type "apropos word" to search for commands related to "word"...
Reading symbols from python...
(gdb) run example.py 
Starting program: /home/tytskiy/miniforge3/bin/python example.py
[Thread debugging using libthread_db enabled]
Using host libthread_db library "/usr/lib/x86_64-linux-gnu/libthread_db.so.1".

Program received signal SIGSEGV, Segmentation fault.
0x00007fffe641d36e in torch::jit::Node::addInput(torch::jit::Value*) () from /home/tytskiy/miniforge3/lib/python3.12/site-packages/torch/lib/libtorch_cpu.so
(gdb) 
```

**Expected**:
Get onnx file with PackedMultiHeadAttention:

[packed_attention.onnx.zip](https://github.com/user-attachments/files/18489000/packed_attention.onnx.zip)

www.netron.app shows it:
<img width="323" alt="Image" src="https://github.com/user-attachments/assets/2a0f8864-687d-4e97-9a0e-abdd97d20c74" />


**Why:**
I guess it happens because we pass None here
https://github.com/pytorch/pytorch/blob/b5655d9821b7214af200d0b8796a10ad34b85229/torch/onnx/_internal/jit_utils.py#L290-L292

and interpret it as Value* here
https://github.com/pytorch/pytorch/blob/803017f3cb73bb115eda5ec0e0a19688ccafbf4e/torch/csrc/jit/python/python_ir.cpp#L585-L589

**Possible fix 1:**
Replace all ```None``` arguments with magic line```g.op('prim::Constant').setType(_C.OptionalType.ofTensor())```. But it seems unclear and it not covered by docs.

**Possible fix 2:**
Just replace this line
https://github.com/pytorch/pytorch/blob/b5655d9821b7214af200d0b8796a10ad34b85229/torch/onnx/_internal/jit_utils.py#L265-L267

with

```python
def _const_if_tensor(graph_context: GraphContext, arg):
    if arg is None:
        return graph_context.op('prim::Constant').setType(_C.OptionalType.ofTensor())
```


### Versions

It seems to happen in any environment

...
Versions of relevant libraries:
[pip3] onnx==1.17.0
[pip3] onnxscript==0.1.0.dev20241203
[pip3] onnxsim==0.4.36
[pip3] torch==2.5.1
[pip3] torchaudio==2.5.1
[pip3] torchvision==0.20.1