Run decomp before processing (pytorch#5713)

test/stablehlo/test_exports.py

@@ -0,0 +1,32 @@
+import unittest
+import torch
+import torch.nn.functional as F
+from torch_xla.stablehlo import exported_program_to_stablehlo
+
+
+class Interpolate(torch.nn.Module):
+
+  def forward(self, masks: torch.Tensor) -> torch.Tensor:
+    masks = F.interpolate(
+        masks,
+        size=(500, 500),
+        mode="bilinear",
+        align_corners=False,
+    )
+    return masks
+
+
+class ExportTest(unittest.TestCase):
+
+  def test_interpolate(self):
+
+    arg = (torch.randn(3, 3, 200, 200),)
+    model = Interpolate()
+
+    ans = model(*arg)
+
+    with torch.no_grad():
+      exported = torch._export.export(model, arg)
+      shlo = exported_program_to_stablehlo(exported)
+      ans2 = shlo(*arg).cpu().to(torch.float32)
+      self.assertTrue(torch.allclose(ans, ans2, atol=1e-5))

torch_xla/stablehlo.py

@@ -240,7 +240,7 @@ def _exported_program_to_stablehlo_bundle(exported_model,
                                           options) -> StableHLOModelBundle:
   if options is None:
     options = StableHLOExportOptions()
-
+  exported_model = exported_model.run_decompositions()
   input_args = _extract_input_args(exported_model, options)
 
   device = xm.xla_device()