Repository: pytorch/torchdynamo

Issue:
TORCHDYNAMO_DYNAMIC_SHAPES=1 expects sizes returned from graph to be tuples, but AOTAutograd doesn't agree
Repro: `TORCH_SHOW_CPP_STACKTRACES=1 TORCHDYNAMO_DYNAMIC_SHAPES=1 AOT_DYNAMIC_SHAPES=1 TORCH_MONKEYPATCH_GETITEM=1 AOT_FX_GRAPHS=1 python benchmarks/torchbench.py --only hf_T5 --accuracy --backend aot_eager --training)` with symbolic-shapes branch https://github.com/pytorch/pytorch/pull/84246

We end up with this AOTAutograd graph:

```
====== Forward graph ======                                                                                                                                
class GraphModule(torch.nn.Module):                                                                                                                        
    def forward(self, primals_1, primals_2, primals_3):                                                                                                    
                                                                                                                                                           
        # Module stack: {}, File: /home/ezyang/local/pytorch-tmp-env/lib/python3.9/site-packages/transformers/models/t5/modeling_t5.py:925, code: input_ids
 = input_ids.view(-1, input_shape[-1])                                                                                                                     
        sym_size = torch.ops.aten.sym_size(primals_2, 1)                                                                                                   
        view = torch.ops.aten.view.default(primals_2, [-1, sym_size])                                                                                      
                                                                             
        # Module stack: {'self_embed_tokens': 'Embedding'}, File: /home/ezyang/local/pytorch-tmp-env/lib/python3.9/site-packages/transformers/models/t5/mod
eling_t5.py:934, code: inputs_embeds = self.embed_tokens(input_ids)          
        embedding = torch.ops.aten.embedding.default(primals_1, view);  primals_1 = None                                                                   
                                                                             
        # Module stack: {}, File: /home/ezyang/local/pytorch-tmp-env/lib/python3.9/site-packages/transformers/models/t5/modeling_t5.py:945, code: attention
_mask = torch.ones(batch_size, mask_seq_length).to(inputs_embeds.device)                                                                                   
        sym_size_1 = torch.ops.aten.sym_size(primals_2, 0);  primals_2 = None                                                                              
        ones = torch.ops.aten.ones.default([sym_size_1, sym_size], device = device(type='cpu'), pin_memory = False)                              
        detach = torch.ops.aten.detach.default(ones);  ones = None                                                                                         
        detach_1 = torch.ops.aten.detach.default(detach);  detach = None                                                                                   
        _to_copy = torch.ops.aten._to_copy.default(detach_1, dtype = torch.float32, layout = torch.strided, device = device(type='cuda', index=0));  detach
_1 = None                                                                    
                                                                                                                                                           
        # Module stack: {}, File: /home/ezyang/local/pytorch-tmp-env/lib/python3.9/site-packages/transformers/models/t5/modeling_t5.py:948, code: encoder_a
ttention_mask = torch.ones(                                                                                                                                
        sym_size_2 = torch.ops.aten.sym_size(primals_3, 1);  primals_3 = None                                                                              
        ones_1 = torch.ops.aten.ones.default([sym_size_1, sym_size_2], dtype = torch.int64, device = device(type='cuda', index=0), pin_memory = False);  sy
m_size_2 = None                                                              
        detach_2 = torch.ops.aten.detach.default(ones_1);  ones_1 = None                                                                                   
        detach_3 = torch.ops.aten.detach.default(detach_2);  detach_2 = None                                                                               
        return [_to_copy, sym_size_1, sym_size, detach_3, embedding, view] 

aot output: (tensor([[1., 1., 1.,  ..., 1., 1., 1.],                                                                                                       
        [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0',                                                                                                  
       grad_fn=<CompiledFunctionBackward>), 2, 1024, tensor([[1, 1, 1,  ..., 1, 1, 1],                                                                     
        [1, 1, 1,  ..., 1, 1, 1]], device='cuda:0'))  
```

but Dynamo is expecting (2, 1024)

Comments from @jansel :

Actually, perhaps we should just change `DynamicShapeVariable(len=2)` to `TupleVariable([DynamicShapeVariable, DynamicShapeVariable])`
There is actually a SizeVariable() in [lists.py](https://lists.py/). This is only a dynamic_shapes=True thing.
(dynamic_shapes=True is less tested)

We live debugged this on the stream at https://www.twitch.tv/edwardzyang


Discussion:
This fix should involve changing these lines:
https://github.com/pytorch/torchdynamo/blob/7350f88e00a4e03c6eb591546727c25151e5ca56/torchdynamo/variables/tensor.py#L180-L187

To return a `torchdynamo.variables.SizeVariable([DynamicShapeVariable(), ...])` (which is a type of TupleVariable).
Actually it would look somewhat similar to the static shapes case right below:
https://github.com/pytorch/torchdynamo/blob/7350f88e00a4e03c6eb591546727c25151e5ca56/torchdynamo/variables/tensor.py#L196-L200
Would need to get the items from the proxy not the example_value.
```
diff --git a/torchdynamo/variables/tensor.py b/torchdynamo/variables/tensor.py
index 0ec864de..c04ef61a 100644
--- a/torchdynamo/variables/tensor.py
+++ b/torchdynamo/variables/tensor.py
@@ -178,13 +178,14 @@ class TensorVariable(VariableTracker):
             options.update(specialized_props)
             return cls(proxy, **options)
         elif (
-            istype(example_value, (torch.Size, int, bool, float))
+            istype(example_value, (int, bool, float))
             and config.dynamic_shapes
         ):
             proxy.node.meta["example_value"] = example_value
-            if isinstance(example_value, torch.Size):
-                options["dyn_shape_len"] = len(example_value)
             return DynamicShapeVariable(proxy, type(example_value), **options)
+        elif istype(example_value, torch.Size) and config.dynamic_shapes:
+            sizes = [DynamicShapeVariable(proxy[i], int) for i in range(len(example_value))]
+            return SizeVariable(sizes, **options)
         elif istype(example_value, int) and proxy.node.target in (
             torch.seed,
             operator.mod,
```

just need to write a test

Resolving patch:
diff --git a/torchdynamo/variables/lists.py b/torchdynamo/variables/lists.py
--- a/torchdynamo/variables/lists.py
+++ b/torchdynamo/variables/lists.py
@@ -1,7 +1,9 @@
 from typing import Dict
 from typing import List
+from typing import Optional
 
 import torch
+import torch.fx
 
 from .. import config
 from .. import variables
@@ -38,6 +40,7 @@ def as_python_constant(self):
         return self.python_type()([x.as_python_constant() for x in self.items])
 
     def as_proxy(self):
+        assert self.python_type() is not SizeVariable
         return self.python_type()(self._as_proxy())
 
     def getitem_const(self, arg: VariableTracker):
@@ -247,9 +250,59 @@ def call_method(
 class SizeVariable(TupleVariable):
     """torch.Size(...)"""
 
+    def __init__(
+        self,
+        items: List[VariableTracker],
+        proxy: Optional[torch.fx.Proxy] = None,
+        **kwargs,
+    ):
+        self.proxy = proxy
+        super().__init__(items, **kwargs)
+
     def python_type(self):
         return torch.Size
 
+    def as_proxy(self):
+        if self.proxy is not None:
+            return self.proxy
+
+        # torch.Size needs special handling.  Normally, we pun a list-like
+        # container to directly contain Proxy/Node objects from FX, and FX
+        # knows to look inside containers (via map_aggregate).  But torch.Size
+        # is weird; although it subclasses from tuple, it doesn't allow
+        # members which aren't int-like (rejecting Proxy and Node).  This
+        # means we can't use the normal representation trick
+        # torch.Size([proxy0, proxy1]).  I looked into seeing if I could
+        # relax torch.Size in PyTorch proper, but if torch.Size constructor
+        # sees a type that it doesn't recognize, it will try to call
+        # __index__() on it, so there is no BC way to actually change this
+        # behavior (though it occurs to me that I could have just added a
+        # YOLO no checking alternate constructor.)
+        #
+        # To work around this problem, I represent a torch.Size proxy as
+        # a straight up proxy, that would have been constructed by taking
+        # the constituent proxies as arguments.  This trick can be generally
+        # used for any construct that we need a proxy for but we can't
+        # directly represent as an aggregate; I don't see very many examples
+        # of this in torchdynamo though!
+
+        # Look for a proxy.  If there are none, do the legacy behavior
+        tracer = None
+        proxies = self._as_proxy()
+        for proxy in proxies:
+            if isinstance(proxy, torch.fx.Proxy):
+                tracer = proxy.tracer
+                break
+
+        if tracer is None:
+            return torch.Size(proxies)
+
+        proxy = tracer.create_proxy("call_function", torch.Size, (proxies,), {})
+        proxy.node.meta["example_value"] = torch.Size(
+            [p.node.meta["example_value"] for p in proxies]
+        )
+        return proxy
+
     def reconstruct(self, codegen):
         codegen.load_import_from("torch", "Size")
         codegen.foreach(self.items)
diff --git a/torchdynamo/variables/tensor.py b/torchdynamo/variables/tensor.py
--- a/torchdynamo/variables/tensor.py
+++ b/torchdynamo/variables/tensor.py
@@ -177,14 +177,17 @@ def context():
 
             options.update(specialized_props)
             return cls(proxy, **options)
-        elif (
-            istype(example_value, (torch.Size, int, bool, float))
-            and config.dynamic_shapes
-        ):
+        elif istype(example_value, (int, bool, float)) and config.dynamic_shapes:
             proxy.node.meta["example_value"] = example_value
-            if isinstance(example_value, torch.Size):
-                options["dyn_shape_len"] = len(example_value)
             return DynamicShapeVariable(proxy, type(example_value), **options)
+        elif istype(example_value, torch.Size) and config.dynamic_shapes:
+            proxy.node.meta["example_value"] = example_value
+            sizes = []
+            for i, v in enumerate(example_value):
+                proxy_i = proxy[i]
+                proxy_i.node.meta["example_value"] = v
+                sizes.append(DynamicShapeVariable(proxy_i, int))
+            return SizeVariable(sizes, proxy, **options)
         elif istype(example_value, int) and proxy.node.target in (
             torch.seed,
             operator.mod,
@@ -502,22 +505,18 @@ def call_method(
 
 
 class DynamicShapeVariable(TensorVariable):
-    def __init__(self, proxy, dyn_shape_cls, dyn_shape_len=None, **kwargs):
+    """
+    Represents a symbolic size, e.g., as returned by tensor.size(0)
+    """
+
+    def __init__(self, proxy, dyn_shape_cls, **kwargs):
         super(DynamicShapeVariable, self).__init__(proxy, **kwargs)
         self.dyn_shape_cls = dyn_shape_cls
-        self.dyn_shape_len = dyn_shape_len
 
     def python_type(self):
         return self.dyn_shape_cls
 
     def unpack_var_sequence(self, tx):
-        if self.dyn_shape_len is not None:
-            return [
-                variables.BuiltinVariable(
-                    operator.getitem, **VariableTracker.propagate(self)
-                ).call_function(tx, [self, variables.ConstantVariable(i)], {})
-                for i in range(self.dyn_shape_len)
-            ]
         super(DynamicShapeVariable, self).unpack_var_sequence(tx)
 
 


Test patch:
diff --git a/test/test_repros.py b/test/test_repros.py
--- a/test/test_repros.py
+++ b/test/test_repros.py
@@ -885,7 +885,7 @@ def test_hf_t5_forward(self):
         self.assertTrue(same(opt_model(input), correct))
 
         self.assertEqual(cnt.frame_count, 1)
-        self.assertEqual(cnt.op_count, ifdyn(14, 11))
+        self.assertEqual(cnt.op_count, ifdyn(13, 11))
 
     def test_slicing_dynamic_shape(self):
         def fn(y):
@@ -1608,6 +1608,25 @@ def fn(x):
         opt_fn(x)
         self.assertEqual(cnt.frame_count, 1)
 
+    # This doesn't work without fake tensors but I don't care
+    @patch.object(torchdynamo.config, "fake_tensor_propagation", True)
+    def test_issue1466_size_aot_autograd(self):
+        def fn(x):
+            # do a tensor op and a size compute
+            y = x * 2
+            x_size = x.size()
+            # trigger a graph break
+            print("arf")
+            # use the tensor op and size compute
+            z = y.view(x_size) + 1
+            return z
+
+        x = torch.randn(2, 3, requires_grad=True)
+        ref = fn(x)
+        opt_fn = torchdynamo.optimize("aot_eager")(fn)
+        res = opt_fn(x)
+        self.assertTrue(same(ref, res))
+
     def test_ellipsis(self):
         class Repro(torch.nn.Module):
             def __init__(self):