Repository: pytorch
Issue: torch.compile Jamba: Scheduler Error in codegen for ComputedBuffer
Labels: triaged,oncall: pt2,module: inductor

### 🐛 Describe the bug

From user-empathy day. Error message:
```
CRITICAL:torch._inductor.scheduler:Error in codegen for ComputedBuffer(name='buf7', layout=MutationLayout('cpu', torch.float32, size=[1, 100, 8], stride=[800, 8, 1]), data=Scatter(device=device(type='cpu'), dtype=tor
ch.float32, inner_fn=<function scatter_reduce_.<locals>.fn at 0x7faf1989cf70>, ranges=[1, 100, 1], output_indexer=<function scatter_reduce_.<locals>.output_indexer at 0x7faf1989e5f0>, scatter_mode=None))
```

repro: https://github.com/xmfan/Jamba/commit/be6b5c12b23862669b37c8ee41cc1f8938529266 (Heads up: 22 minutes to compile for cpu inductor)
```
Warm up iteration 0                                                                                                                                                                                                     
CRITICAL:torch._inductor.scheduler:Error in codegen for ComputedBuffer(name='buf7', layout=MutationLayout('cpu', torch.float32, size=[1, 100, 8], stride=[800, 8, 1]), data=Scatter(device=device(type='cpu'), dtype=tor
ch.float32, inner_fn=<function scatter_reduce_.<locals>.fn at 0x7faf1989cf70>, ranges=[1, 100, 1], output_indexer=<function scatter_reduce_.<locals>.output_indexer at 0x7faf1989e5f0>, scatter_mode=None))             
CRITICAL:torch._inductor.scheduler:Error in codegen for ComputedBuffer(name='buf7', layout=MutationLayout('cpu', torch.float32, size=[1, 100, 8], stride=[800, 8, 1]), data=Scatter(device=device(type='cpu'), dtype=tor
ch.float32, inner_fn=<function scatter_reduce_.<locals>.fn at 0x7fae331fc3a0>, ranges=[1, 100, 1], output_indexer=<function scatter_reduce_.<locals>.output_indexer at 0x7fae331fe440>, scatter_mode=None))             
CRITICAL:torch._inductor.scheduler:Error in codegen for ComputedBuffer(name='buf7', layout=MutationLayout('cpu', torch.float32, size=[1, 100, 8], stride=[800, 8, 1]), data=Scatter(device=device(type='cpu'), dtype=tor
ch.float32, inner_fn=<function scatter_reduce_.<locals>.fn at 0x7faf13858550>, ranges=[1, 100, 1], output_indexer=<function scatter_reduce_.<locals>.output_indexer at 0x7faf1385b910>, scatter_mode=None))             
CRITICAL:torch._inductor.scheduler:Error in codegen for ComputedBuffer(name='buf7', layout=MutationLayout('cpu', torch.float32, size=[1, 100, 8], stride=[800, 8, 1]), data=Scatter(device=device(type='cpu'), dtype=tor
ch.float32, inner_fn=<function scatter_reduce_.<locals>.fn at 0x7fae359348b0>, ranges=[1, 100, 1], output_indexer=<function scatter_reduce_.<locals>.output_indexer at 0x7fae49d80d30>, scatter_mode=None))             
CRITICAL:torch._inductor.scheduler:Error in codegen for ComputedBuffer(name='buf7', layout=MutationLayout('cpu', torch.float32, size=[1, 100, 8], stride=[800, 8, 1]), data=Scatter(device=device(type='cpu'), dtype=tor                                                                                                          
ch.float32, inner_fn=<function scatter_reduce_.<locals>.fn at 0x7fae490a4550>, ranges=[1, 100, 1], output_indexer=<function scatter_reduce_.<locals>.output_indexer at 0x7fae490a60e0>, scatter_mode=None))                                                                                                                       
CRITICAL:torch._inductor.scheduler:Error in codegen for ComputedBuffer(name='buf7', layout=MutationLayout('cpu', torch.float32, size=[1, 100, 8], stride=[800, 8, 1]), data=Scatter(device=device(type='cpu'), dtype=tor                                                                                                          
ch.float32, inner_fn=<function scatter_reduce_.<locals>.fn at 0x7fae49027520>, ranges=[1, 100, 1], output_indexer=<function scatter_reduce_.<locals>.output_indexer at 0x7fae4956beb0>, scatter_mode=None))             
Took 1355.3630783557892 seconds                                                                                                                                                                                                                                                                                                   
defaultdict(<class 'collections.Counter'>, {'frames': Counter({'total': 67, 'ok': 65}), 'inline_call': Counter({'data dependent operator: aten._local_scalar_dense.default; to enable, set torch._dynamo.config.capture_
scalar_outputs = True': 45}), 'unimplemented': Counter({'cache_size_limit reached': 1, "Guard setup for uninitialized class <class 'zeta.nn.modules.multi_input_multi_output.OutputHead'>": 1}), 'graph_break': Counter(                                                                                                          
...
Warm up iteration 1                                                                                                                                                                                                     
Took 17.7934730052948 seconds 
```

Although this is a CRITICAL log, the run completes

### Versions

2.3.0 (fixed by zeta)

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire