Repository: pytorch
Issue: Inconsistent Behavior of `torch.dsplit` with torch.compile
Labels: low priority,triaged,bug,oncall: pt2,module: dynamo,topic: fuzzer,internal ramp-up task

### 🐛 Describe the bug

When running in torch.compile mode, the following error is encountered:

```
**{'indices': [1, 1, 1, 1], 'input': FakeTensor(..., size=(1, 1))}):
torch.dsplit requires a tensor with at least 3 dimensions, but got a tensor with 2 dimensions!**
```
This error occurs despite providing a 6-dimensional tensor as input to torch.dsplit. Furthermore, **disabling the line** _ = torch.special.logit(eps=0.2518168952606683, input=inp2, out=inp1) in the model results in the error not occurring.

### Expected Behavior
torch.dsplit should process the 6-dimensional tensor without error in both eager and torch.compile modes, and the behavior should not be affected by the presence of a torch.special.logit call.

### Error logs
```
torch version:  2.3.0.dev20240130+cu118
==== Eager mode ====
inp1 size : torch.Size([1, 1, 1, 1, 1, 1])
prog.py:13: UserWarning: An output with one or more elements was resized since it had shape [1, 1, 1, 1, 1, 1], which does not match the required output shape [1, 1]. This behavior is deprecated, and in a future PyTorch release outputs will not be resized unless they have zero elements. You can explicitly reuse an out tensor t by resizing it, inplace, to zero elements with t.resize_(0). (Triggered internally at ../aten/src/ATen/native/Resize.cpp:28.)
  _ = torch.special.logit(eps=0.2518168952606683, input=inp2, out=inp1)
==== TorchComp mode ====
inp1 size : torch.Size([1, 1])
Traceback (most recent call last):
  File "prog.py", line 21, in <module>
    ret_exported = torch.compile(model)(*inputs)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py", line 452, in _fn
    return fn(*args, **kwargs)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "prog.py", line 11, in forward
    print("inp1 size :", inp1.size())
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py", line 614, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 748, in _convert_frame
    result = inner_convert(frame, cache_entry, hooks, frame_state)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 390, in _convert_frame_assert
    return _compile(
  File "/opt/conda/envs/std/lib/python3.8/contextlib.py", line 75, in inner
    return func(*args, **kwds)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 650, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/utils.py", line 248, in time_wrapper
    r = func(*args, **kwargs)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 531, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/bytecode_transformation.py", line 1033, in transform_code_object
    transformations(instructions, code_options)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 155, in _fn
    return fn(*args, **kwargs)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 496, in transform
    tracer.run()
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py", line 2125, in run
    super().run()
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py", line 787, in run
    and self.step()
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py", line 750, in step
    getattr(self, inst.opname)(inst)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py", line 469, in wrapper
    return inner_fn(self, inst)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py", line 1249, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py", line 651, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/variables/torch.py", line 614, in call_function
    tensor_variable = wrap_fx_proxy(
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/variables/builder.py", line 1285, in wrap_fx_proxy
    return wrap_fx_proxy_cls(target_cls=TensorVariable, **kwargs)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/variables/builder.py", line 1370, in wrap_fx_proxy_cls
    example_value = get_fake_value(proxy.node, tx, allow_non_graph_fake=True)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/utils.py", line 1653, in get_fake_value
    raise TorchRuntimeError(str(e)).with_traceback(e.__traceback__) from None
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/utils.py", line 1599, in get_fake_value
    ret_val = wrap_fake_exception(
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/utils.py", line 1140, in wrap_fake_exception
    return fn()
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/utils.py", line 1600, in <lambda>
    lambda: run_node(tx.output, node, args, kwargs, nnmodule)
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/utils.py", line 1720, in run_node
    raise RuntimeError(fn_str + str(e)).with_traceback(e.__traceback__) from e
  File "/opt/conda/envs/std/lib/python3.8/site-packages/torch/_dynamo/utils.py", line 1699, in run_node
    return node.target(*args, **kwargs)
torch._dynamo.exc.TorchRuntimeError: Failed running call_function <built-in method dsplit of type object at 0x7fe66c46efa0>(*(), **{'indices': [1, 1, 1, 1], 'input': FakeTensor(..., size=(1, 1))}):
torch.dsplit requires a tensor with at least 3 dimension, but got a tensor with 2 dimensions!

from user code:
   File "prog.py", line 12, in torch_dynamo_resume_in_forward_at_11
    v9_0, _, _, _, _ = torch.dsplit(indices=[1, 1, 1, 1], input=inp1)

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True
```
### Minified repro
```python
import torch
import torch.nn as nn
print("torch version: ",torch.__version__)

class Model(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, inp1, inp2):
        print("inp1 size :", inp1.size())
        v9_0, _, _, _, _ = torch.dsplit(indices=[1, 1, 1, 1], input=inp1)
        _ = torch.special.logit(eps=0.2518168952606683, input=inp2, out=inp1)
        return v9_0, 
inputs = [ torch.randn(1, 1, 1, 1, 1, 1), torch.randn(1, 1) ]
model = Model().to(torch.device("cpu"))
print('==== Eager mode ====')
ret_eager = model(*inputs)

print('==== TorchComp mode ====')
ret_exported = torch.compile(model)(*inputs)
print('OK!')
```

### Versions

PyTorch version: 2.3.0.dev20240130+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.8.17 (default, Jul  5 2023, 21:04:15)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-6.2.0-39-generic-x86_64-with-glibc2.17
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
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:                      48 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             256
On-line CPU(s) list:                0-255
Vendor ID:                          AuthenticAMD
Model name:                         AMD EPYC 7763 64-Core Processor
CPU family:                         25
Model:                              1
Thread(s) per core:                 2
Core(s) per socket:                 64
Socket(s):                          2
Stepping:                           1
Frequency boost:                    enabled
CPU max MHz:                        3529.0520
CPU min MHz:                        1500.0000
BogoMIPS:                           4890.70
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin brs arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif v_spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca fsrm
Virtualization:                     AMD-V
L1d cache:                          4 MiB (128 instances)
L1i cache:                          4 MiB (128 instances)
L2 cache:                           64 MiB (128 instances)
L3 cache:                           512 MiB (16 instances)
NUMA node(s):                       2
NUMA node0 CPU(s):                  0-63,128-191
NUMA node1 CPU(s):                  64-127,192-255
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Mitigation; safe RET
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] pytorch-triton==3.0.0+901819d2b6
[pip3] torch==2.3.0.dev20240130+cu118
[pip3] torchaudio==2.2.0.dev20240130+cu118
[pip3] torchvision==0.18.0.dev20240130+cu118
[conda] numpy                     1.23.5                   pypi_0    pypi
[conda] pytorch-triton            3.0.0+901819d2b6          pypi_0    pypi
[conda] torch                     2.3.0.dev20240130+cu118          pypi_0    pypi
[conda] torchaudio                2.2.0.dev20240130+cu118          pypi_0    pypi
[conda] torchvision               0.18.0.dev20240130+cu118          pypi_0    pypi

cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames @rec @msaroufim @bdhirsh @anijain2305 @zou3519