Repository: pytorch
Issue: [torch.compile] Dynamic shape behavior is different between using torch.compile with and without compiled_autograd.enable
Labels: triaged,oncall: pt2,module: dynamic shapes,module: compiled autograd,module: pt2-dispatcher

### 🐛 Describe the bug

When the model is compiled using torch.compile the backward graph has dynamic shaped tensors but the same model when wrapped with compiled_autograd.enable() the backward graph has static shapes.  The behavior is different between both cases.  Is there any specific reason for this difference in behavior?

### Error logs

With compiled_autograd.enable() example inputs
FakeTensor(..., size=())
FakeTensor(..., size=(4, 10))
FakeTensor(..., size=(4, 5))
FakeTensor(..., size=(4, 50))
FakeTensor(..., size=(50, 50))
FakeTensor(..., size=(4, 50))
FakeTensor(..., size=(50, 5))
FakeTensor(..., size=(4, 5))
Parameter(FakeTensor(..., size=(5,), requires_grad=True))
Parameter(FakeTensor(..., size=(5, 50), requires_grad=True))
Parameter(FakeTensor(..., size=(50,), requires_grad=True))
Parameter(FakeTensor(..., size=(50, 50), requires_grad=True))
Parameter(FakeTensor(..., size=(50,), requires_grad=True))
Parameter(FakeTensor(..., size=(50, 10), requires_grad=True))

Same model without compiled_autograd.enable()
FakeTensor(..., size=(s0, 10))
FakeTensor(..., size=(s0, 5))
FakeTensor(..., size=(s0, 50))
FakeTensor(..., size=(50, 50))
FakeTensor(..., size=(s0, 50))
FakeTensor(..., size=(50, 5))
FakeTensor(..., size=(s0, 5))
FakeTensor(..., size=())


### Minified repro

```python
import torch
import os
import torch.nn as nn
import torch._dynamo
from typing import List
from torch._dynamo.backends.registry import register_backend
from torch._dynamo import compiled_autograd
from torch._dynamo.backends.common import aot_autograd

def custom_inner_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
    print("custom backend called")
    for input in example_inputs:
        print(input)
    #print(gm.graph)
    return gm.forward

@register_backend
def custom_backend(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
    return aot_autograd(
            fw_compiler=custom_inner_compiler,
            bw_compiler=custom_inner_compiler,
            )(gm, example_inputs)

def compiler_fn(gm):
    return torch.compile(gm, backend="custom_backend")

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        m.bias.data.fill_(0.01)

class ToyModel2(nn.Module):
    def __init__(self, in_feat=10, hidden_feat=50, out_feat=5):
        super().__init__()
        self.linear1 = nn.Linear(in_feat, hidden_feat)
        self.linear2 = nn.Linear(hidden_feat, hidden_feat)
        self.linear3 = nn.Linear(hidden_feat, out_feat)
        self.mse_loss = torch.nn.MSELoss()

    def forward(self, inputs, output):
        out1 = self.linear1(inputs)
        out2 = self.linear2(out1)
        out3 = self.linear3(out2)
        return self.mse_loss(out3, output)

def run_single_node():
    device = torch.device("cpu")
    m = ToyModel2()
    m.apply(init_weights)
    m.to(device)
    m = torch.compile(m, backend="custom_backend")
    if True:
    #with compiled_autograd.enable(compiler_fn):
        for cnt in range(2):
            if cnt == 0:
               inp = torch.randn(3,10).to(torch.float32).to(device)
               target = torch.randn(3,5).to(torch.float32).to(device)
            else:
               inp = torch.randn(4,10).to(torch.float32).to(device)
               target = torch.randn(4,5).to(torch.float32).to(device)
            loss = m(inp, target)
            loss.backward()

run_single_node()
```
### Versions

Collecting environment information...
PyTorch version: 2.2.0.dev20231106+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.26.4
Libc version: glibc-2.31

Python version: 3.10.9 (main, Jan 11 2023, 15:21:40) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.4.0-163-generic-x86_64-with-glibc2.31
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
Byte Order:                         Little Endian
Address sizes:                      43 bits physical, 48 bits virtual
CPU(s):                             12
On-line CPU(s) list:                0-11
Thread(s) per core:                 1
Core(s) per socket:                 1
Socket(s):                          12
NUMA node(s):                       1
Vendor ID:                          GenuineIntel
CPU family:                         6
Model:                              85
Model name:                         Intel(R) Xeon(R) Gold 5220R CPU @ 2.20GHz
Stepping:                           0
CPU MHz:                            2194.843
BogoMIPS:                           4389.68
Virtualization:                     VT-x
Hypervisor vendor:                  VMware
Virtualization type:                full
L1d cache:                          384 KiB
L1i cache:                          384 KiB
L2 cache:                           12 MiB
L3 cache:                           429 MiB
NUMA node0 CPU(s):                  0-11
Vulnerability Gather data sampling: Unknown: Dependent on hypervisor status
Vulnerability Itlb multihit:        KVM: Mitigation: Split huge pages
Vulnerability L1tf:                 Mitigation; PTE Inversion; VMX flush not necessary, SMT disabled
Vulnerability Mds:                  Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Meltdown:             Mitigation; PTI
Vulnerability Mmio stale data:      Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Retbleed:             Mitigation; IBRS
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; IBRS, IBPB conditional, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon nopl xtopology tsc_reliable nonstop_tsc cpuid pni pclmulqdq vmx ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single pti ssbd ibrs ibpb stibp tpr_shadow vnmi ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 invpcid avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xsaves arat pku ospke md_clear flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] torch==2.2.0.dev20231106+cpu
[pip3] torchaudio==2.2.0.dev20231106+cpu
[pip3] torchvision==0.17.0.dev20231106+cpu
[conda] numpy                     1.24.3                   pypi_0    pypi
[conda] torch                     2.2.0.dev20231106+cpu          pypi_0    pypi
[conda] torchaudio                2.2.0.dev20231106+cpu          pypi_0    pypi
[conda] torchvision               0.17.0.dev20231106+cpu          pypi_0    pypi


cc @ezyang @chauhang @penguinwu @xmfan @zou3519 @bdhirsh @msaroufim @wconstab @anijain2305