Repository: pytorch
Issue: Mark Dynamic does not work for nn module constructor inputs
Labels: triaged,module: dynamic shapes

for the following code if we compile it we get the following graph :
code:
```
import torch
import torch._dynamo.config
from torch.utils.checkpoint import _is_compiling

@torch.compile()

class Y(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x.view(-1,  self._N_input, 192) 

@torch.compile(dynamic=True)
class M(torch.nn.Module):
    def __init__(self, input):
        super().__init__()
        self._N_input = input.size()[0]

    def forward(self, x, z):
         # input is [B, n * d]
        x = x*2
        x = x*5
        x = x *3
        y = Y()        
        y._N_input = self._N_input
        x =  y.forward(x)# [B, n, d]
        x = x*20
        x = x*30 
        x = x*43
        return x
x = torch.randn(5, 3210, 192)  # Shape [B=4, n*d=12]
num_inputs = torch.randn(3210) 
m = M(num_inputs)
y1, z1 = 3210, 192  # y1 * z1 == 12
output1 = m(x, z1)
print(output1)
```

graph: Note L_self_N_input is Sym(s1) 
```
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code] TRACED GRAPH
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]  ===== __compiled_fn_1 =====
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]  /home/lsakka/pytorch/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]     def forward(self, s0: "Sym(s0)", s1: "Sym(s1)", L_x_: "f32[s0, s1, 192][192*s1, 192, 1]cpu", L_self_N_input: "Sym(s1)"):
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         l_x_ = L_x_
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         l_self_n_input = L_self_N_input
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:751 in forward, code: x = x*2
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         x: "f32[s0, s1, 192][192*s1, 192, 1]cpu" = l_x_ * 2;  l_x_ = None
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:752 in forward, code: x = x*5
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         x_1: "f32[s0, s1, 192][192*s1, 192, 1]cpu" = x * 5;  x = None
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:753 in forward, code: x = x *3
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         x_2: "f32[s0, s1, 192][192*s1, 192, 1]cpu" = x_1 * 3;  x_1 = None
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:737 in __init__, code: super().__init__()
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         _log_api_usage_once = torch._C._log_api_usage_once('python.nn_module');  _log_api_usage_once = None
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:740 in forward, code: return x.view(-1,  self._N_input, 192)
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         x_3: "f32[s0, s1, 192][192*s1, 192, 1]cpu" = x_2.view(-1, l_self_n_input, 192);  x_2 = l_self_n_input = None
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:759 in forward, code: x = x*20
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         x_4: "f32[s0, s1, 192][192*s1, 192, 1]cpu" = x_3 * 20;  x_3 = None
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:760 in forward, code: x = x*30
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         x_5: "f32[s0, s1, 192][192*s1, 192, 1]cpu" = x_4 * 30;  x_4 = None
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:761 in forward, code: x = x*43
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         x_6: "f32[s0, s1, 192][192*s1, 192, 1]cpu" = x_5 * 43;  x_5 = None
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         return (x_6,)
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:39:29.878000 2775135 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code] 
```

Now if we want to use mark_dynamic instead, I : 
1) remove(dynamic=True)
2) call mark dynamic as the following:
```
mport torch
import torch._dynamo.config
from torch.utils.checkpoint import _is_compiling

@torch.compile()

class Y(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x.view(-1,  self._N_input, 192) 

@torch.compile()
class M(torch.nn.Module):
    def __init__(self, input):
        super().__init__()

        self._N_input = input.size()[0]

    def forward(self, x, z):
         # input is [B, n * d]
        x = x*2
        x = x*5
        x = x *3
        y = Y()        
        y._N_input = self._N_input
        x =  y.forward(x)# [B, n, d]
        x = x*20
        x = x*30 
        x = x*43
        return x


x = torch.randn(5, 3210, 192)  # Shape [B=4, n*d=12]
num_inputs = torch.randn(3210) 

torch._dynamo.decorators.mark_dynamic(x, 1)
torch._dynamo.decorators.mark_dynamic(num_inputs, 0)
m = M(num_inputs)
y1, z1 = 3210, 192  # y1 * z1 == 12
output1 = m(x, z1)
print(output1)
```

but this generate this graph,  looks like mark_dynamic did not work on num_inputs
    raise ConstraintViolationError(
torch.fx.experimental.symbolic_shapes.ConstraintViolationError: Constraints violated (L['x'].size()[1])! For more information, run with TORCH_LOGS="+dynamic".
  - Not all values of RelaxedUnspecConstraint(L['x'].size()[1]) are valid because L['x'].size()[1] was inferred to be a constant (3210).

```
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code] TRACED GRAPH
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]  ===== __compiled_fn_1 =====
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]  /home/lsakka/pytorch/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]     def forward(self, L_x_: "f32[5, 3210, 192][616320, 192, 1]cpu"):
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         l_x_ = L_x_
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:751 in forward, code: x = x*2
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         x: "f32[5, 3210, 192][616320, 192, 1]cpu" = l_x_ * 2;  l_x_ = None
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:752 in forward, code: x = x*5
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         x_1: "f32[5, 3210, 192][616320, 192, 1]cpu" = x * 5;  x = None
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:753 in forward, code: x = x *3
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         x_2: "f32[5, 3210, 192][616320, 192, 1]cpu" = x_1 * 3;  x_1 = None
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:737 in __init__, code: super().__init__()
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         _log_api_usage_once = torch._C._log_api_usage_once('python.nn_module');  _log_api_usage_once = None
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:740 in forward, code: return x.view(-1,  self._N_input, 192)
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         x_3: "f32[5, 3210, 192][616320, 192, 1]cpu" = x_2.view(-1, 3210, 192);  x_2 = None
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:759 in forward, code: x = x*20
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         x_4: "f32[5, 3210, 192][616320, 192, 1]cpu" = x_3 * 20;  x_3 = None
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:760 in forward, code: x = x*30
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         x_5: "f32[5, 3210, 192][616320, 192, 1]cpu" = x_4 * 30;  x_4 = None
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]          # File: /home/lsakka/pytorch/example.py:761 in forward, code: x = x*43
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         x_6: "f32[5, 3210, 192][616320, 192, 1]cpu" = x_5 * 43;  x_5 = None
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         return (x_6,)
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code]         
V0122 22:42:15.613000 2829692 torch/_dynamo/output_graph.py:1353] [0/0] [__graph_code] 
```

cc @chauhang @penguinwu @ezyang @bobrenjc93