Improve getitem syntax for TensorType (pytorch#84555)

Allows `TensorType[Dyn, 3, Dyn]` instead of the prior `TensorType[(Dyn, 3, Dyn)]`.
Pull Request resolved: pytorch#84555
Approved by: https://github.com/jamesr66a

Author: jansel

test/fx/test_gradual_type.py

@@ -44,13 +44,16 @@ def test_annotations(self):
         where n is the corresoinding node in the resulting graph.
         """
         class M(torch.nn.Module):
-            def forward(self, x: TensorType((1, 2, 3, Dyn)), y: Dyn):
-                return torch.add(x, y)
+            def forward(self,
+                        x: TensorType((1, 2, 3, Dyn)),
+                        y: Dyn,
+                        z: TensorType[Dyn, 3, Dyn]):
+                return torch.add(x, y) + z
 
         module = M()
         symbolic_traced: torch.fx.GraphModule = symbolic_trace(module)
 
-        expected_ph_types = [TensorType((1, 2, 3, Dyn)), Dyn]
+        expected_ph_types = [TensorType((1, 2, 3, Dyn)), Dyn, TensorType((Dyn, 3, Dyn))]
         expected_iter = iter(expected_ph_types)
 
         for n in symbolic_traced.graph.nodes:

torch/fx/tensor_type.py

@@ -28,7 +28,9 @@ def __eq__(self, other):
 
     @staticmethod
     def __class_getitem__(*args):
-        return TensorType(args[0])
+        if len(args) == 1 and isinstance(args[0], tuple):
+            args = args[0]
+        return TensorType(tuple(args))
 
 
 class _DynType: