Repository: pytorch
Issue: [Functionalization] `index_reduce_` op tests with functionalization enabled
Labels: triaged,module: meta tensors,module: functionalization

### 🐛 Describe the bug

With functionalization, the existing `index_reduce_` python op test (https://github.com/pytorch/pytorch/blob/master/test/test_torch.py#L3035) fails.

To reproduce (this is one of the sample inputs to the `index_reduce_` test linked above):
```
import torch
import functorch
def test():
	dest = torch.tensor([[[ 0.0322,  1.2734, -3.4688,  8.1875, -4.2500],
	         [-5.6250,  1.3828,  7.7188,  0.3887,  5.5312],
	         [ 3.2344,  5.0312, -7.4062,  1.2422, -0.1719],
	         [-6.0312,  6.2188, -1.1641, -0.3203,  0.2637]],

	        [[-6.9688, -3.5938,  2.6406,  4.3125,  0.1348],
	         [ 4.5000, -0.5938, -5.5312, -1.8281,  1.1562],
	         [ 1.5781, -1.7891,  3.8906,  1.2969,  1.9688],
	         [-6.5000,  2.4375, -4.8125,  3.0312,  1.9453]],

	        [[ 0.2002,  7.7188,  1.5547, -7.6875, -2.5781],
	         [-4.1562,  1.8125,  6.5625,  8.2500,  5.4062],
	         [ 4.2812,  6.5625, -3.3906,  1.7266,  8.8750],
	         [-6.9375,  7.0625,  3.4844, -7.9375,  8.5625]]], dtype=torch.bfloat16)
	idx = torch.tensor([], dtype=torch.int64)
	src = torch.empty((3, 4, 0), dtype=torch.bfloat16)
	dest.index_reduce_(2, idx, src, 'mean', include_self=False)
```
Output:
```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/opt/conda/lib/python3.8/site-packages/torch/_functorch/vmap.py", line 39, in fn
    return f(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/_functorch/eager_transforms.py", line 1582, in wrapped
    func_outputs = func(*func_args, **func_kwargs)
  File "<stdin>", line 18, in test
  File "/opt/conda/lib/python3.8/site-packages/torch/_decomp/decompositions.py", line 3314, in inplace_op
    out = outplace_op(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/_ops.py", line 499, in __call__
    return self._op(*args, **kwargs or {})
IndexError: select(): index 0 out of range for tensor of size [3, 4, 0] at dimension 2
```
Full dispatch trace logs: https://gist.github.com/wonjoolee95/d6c2c31df8a3342ddbf56523c0eeab66
Full dispatch trace logs without functionalization: https://gist.github.com/wonjoolee95/8a9c2543b0b017f9df049da57fc84dce

The error itself seems to be clear -- due to some index out-bound-error, as the code tries to access index 0 at dimension 2 of shape [3, 4, 0] as mentioned in the error logs; however, this only happens when functionalization is enabled. The last few bits of the dispatch trace seems suspicious. Without functionalization, the dispatch looks like:
```
 [call] op=[aten::index_reduce_], key=[AutogradCPU]
  [redispatch] op=[aten::index_reduce_], key=[ADInplaceOrView]
   [redispatch] op=[aten::index_reduce_], key=[CPU]
    [call] op=[aten::to.dtype], key=[CPU]
    [call] op=[aten::index_fill_.int_Scalar], key=[CPU]
     [call] op=[aten::as_strided], key=[CPU]
     [call] op=[aten::as_strided], key=[CPU]
```
However, with functionalization, it looks like:
```
 [call] op=[aten::index_reduce_], key=[FuncTorchDynamicLayerFrontMode]
  [callBoxed] op=[aten::index_reduce_], key=[Functionalize]
   [call] op=[aten::index_reduce_], key=[Meta]
    [callBoxed] op=[aten::index_reduce], key=[Meta]
     [call] op=[aten::select.int], key=[Meta]
      [call] op=[aten::as_strided], key=[Meta]
     [call] op=[aten::select.int], key=[Meta]
```
Just looking at it at a high-level, seems like functionalization now decomposes into `select.int` that might deal with indices differently compared to the previous ops?

Please let me know if you need any more information.

cc @ezyang @eellison @bdhirsh @soumith @alanwaketan 

### Versions

Nightly