Make sure dynamo give us different graph when input shape changes (#5177)

* Make sure dynamo give us different graph when input shape changes

* fix typo

Author: JackCaoG

test/dynamo/test_dynamo.py

@@ -68,6 +68,25 @@ def test_simple_model(self):
     res_cpu_3 = self.fn_simple(x + y, y * 3)
     self.assertTrue(torch.allclose(res_cpu_3, res_xla_dynamo_3.cpu()))
 
+  def test_simple_model_with_different_input_shape(self):
+    met.clear_counters()
+    device = xm.xla_device()
+    xla_x = torch.randn(5, 5).to(device)
+    xla_y = torch.randn(5, 5).to(device)
+    xla_z = torch.randn(10, 10).to(device)
+    self.fn_simple_dynamo(xla_x, xla_x)
+    compile_count = met.metric_data('CompileTime')[0]
+    # Execute with input with same shape should not trigger additional compilation
+    self.fn_simple_dynamo(xla_y, xla_y)
+    self.assertEqual(met.metric_data('CompileTime')[0], compile_count)
+    # Give `fn_simple_dynamo` an input with different shappe, we expect
+    # dynamo to recognize this is a different graph and let XLA to retrace/recompile
+    res_xla_dynamo_3 = self.fn_simple_dynamo(xla_z, xla_z)
+    self.assertEqual(met.metric_data('CompileTime')[0], compile_count + 1)
+    self.assertTrue(
+        torch.allclose(res_xla_dynamo_3.cpu(),
+                       self.fn_simple(xla_z.cpu(), xla_z.cpu())))
+
   def test_resnet18(self):
     device = xm.xla_device()
     batch_size = xu.getenv_as('BATCH_SIZE', int, defval=4)