Make a dynamo inference test that runs on TPU

test/dynamo/test_dynamo.py

@@ -333,19 +333,23 @@ def test_simple_model_with_different_input_shape(self, initialize_on_cuda):
             rtol=1e-05,
             atol=1e-05))
 
+  def get_loader(self, device, sample_count):
+    batch_size = xu.getenv_as('BATCH_SIZE', int, defval=4)
+    loader = xu.SampleGenerator(
+        data=(torch.randn(batch_size, 3, 224, 224, device=device),
+              torch.zeros(batch_size, dtype=torch.int64, device=device)),
+        sample_count=sample_count)
+    return loader
+
   @skipOnTpu
   @parameterized.parameters(
       True,
       False,
   )
   def test_resnet18(self, initialize_on_cuda):
     device = self._choose_proper_device(initialize_on_cuda)
-    batch_size = xu.getenv_as('BATCH_SIZE', int, defval=4)
     sample_count = xu.getenv_as('SAMPLE_COUNT', int, defval=10)
-    loader = xu.SampleGenerator(
-        data=(torch.randn(batch_size, 3, 224, 224, device=device),
-              torch.zeros(batch_size, dtype=torch.int64, device=device)),
-        sample_count=sample_count)
+    loader = self.get_loader(device, sample_count)
     resnet18 = torchvision.models.resnet18()
     resnet18.eval()
     device_resnet18 = torchvision.models.resnet18()
@@ -370,6 +374,32 @@ def test_resnet18(self, initialize_on_cuda):
     self.assertEqual(
         met.metric_data('RunCachedGraphOutputData')[0], sample_count)
 
+  def test_resnet18_lazy_vs_dynamo(self):
+    sample_count = xu.getenv_as('SAMPLE_COUNT', int, defval=10)
+    device = torch_xla.device()
+    loader = self.get_loader(device, sample_count)
+    resnet18_base = torchvision.models.resnet18()
+    resnet18_base.eval()
+    xla_resnet18 = torchvision.models.resnet18()
+    xla_resnet18.load_state_dict(resnet18_base.state_dict())
+    xla_resnet18.to(device)
+    xla_resnet18.eval()
+    resnet18_base.to(device)
+    # materalize the fake data for test purpose
+    xm.mark_step()
+    xm.wait_device_ops()
+    met.clear_all()
+    dynamo_resnet18 = torch.compile(xla_resnet18, backend='openxla')
+    for data, _ in loader:
+      output_lazy = resnet18_base(data)
+      torch_xla.sync()
+      output_dynamo = dynamo_resnet18(data)
+      self.assertTrue(
+          torch.allclose(
+              output_lazy.cpu(), output_dynamo.cpu(), rtol=1e-05, atol=1e-05))
+      # skip the counter/metrics check since LTC also runs on device and will
+      # mess up the counter check.
+
 
 class DynamoCpuFallbackTest(unittest.TestCase):