Repository: pytorch
Issue: torch compile does not work with torch.nn.functional.softmax ?
Labels: triaged,oncall: pt2,module: inductor

### 🐛 Describe the bug

Hi,

I'm trying to compile our model with torch.compile stack.
It looks like its not compatible with softmax.

Here is a full error message.
```
rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING]   File "/mnt/xarfuse/uid-194481/e0143f7f-seed-bac47847-3826-49af-9334
-48541a87a0c7-ns-4026535691/torch/_subclasses/fake_tensor.py", line 1382, in dispatch
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING]     return decomposition_table[func](*args, **kwargs)
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING]   File "/mnt/xarfuse/uid-194481/e0143f7f-seed-bac47847-3826-49af-9334
-48541a87a0c7-ns-4026535691/torch/_refs/__init__.py", line 2309, in amax
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING]     return _reduction(
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING]   File "/mnt/xarfuse/uid-194481/e0143f7f-seed-bac47847-3826-49af-9334
-48541a87a0c7-ns-4026535691/torch/_refs/__init__.py", line 2097, in _reduction
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING]     raise RuntimeError(
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING] torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING] RuntimeError: reducing over zero-size dimension for reduction operation
 without identity
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING] 
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING] While executing %amax : [num_users=1] = call_function[target=torch.ops.aten
.amax.default](args = (%view_4, [-1], True), kwargs = {})
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING] Original traceback:
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING]   File "/mnt/xarfuse/uid-194481/e0143f7f-seed-bac47847-3826-49af-9334
-48541a87a0c7-ns-4026535691/pyspeech/modules/emformer_attention.py", line 955, in <resume in forward>
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING]     rc_output_memory = self.prepare_attention_output(
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING]   File "/mnt/xarfuse/uid-194481/e0143f7f-seed-bac47847-3826-49af-9334
-48541a87a0c7-ns-4026535691/pyspeech/modules/emformer_attention.py", line 811, in prepare_attention_output
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING]     attention_weights_float = torch.nn.functional.softmax(
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING] 
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING] 
[rank11]:[2023-07-28 04:35:51,728] torch._dynamo.convert_frame: [WARNING] 
[rank11]:[2023-07-28 04:35:51,729] torch._dynamo.convert_frame: [WARNING] converting frame raised error, suppressing error

```

Can you help me to look?

### Versions

Meta internal

cc @ezyang @msaroufim @bdhirsh @anijain2305 @zou3519 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @amjames @desertfire @wconstab