[torch_xla2] Fix TF SavedModel export, add test (#8299)

.github/workflows/torch_xla2.yml

@@ -39,6 +39,7 @@ jobs:
         run: |
           pip install -r test-requirements.txt
           pip install -e .[cpu]
+          pip install tensorflow-cpu  # for TF integrations tests
       - name: Run tests
         working-directory: experimental/torch_xla2
         shell: bash

experimental/torch_xla2/test/test_tf_integration.py

@@ -0,0 +1,52 @@
+import jax
+import os
+import tempfile
+import tensorflow as tf
+import torch
+import torch.nn.functional as F
+import torch_xla2
+
+from torch_xla2 import tf_integration
+from . import test_base
+
+
+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 TfIntegrationTest(test_base.TestCase):
+
+  def setUp(self):
+    torch.manual_seed(0)
+    torch_xla2.enable_accuracy_mode()
+
+  def test_interpolate(self):
+    """Simple model roundtripped through TF savedmodel"""
+
+    # Create model
+    arg = (torch.randn(3, 3, 200, 200),)
+    pt_model = Interpolate()
+
+    # Export to SavedModel
+    with tempfile.TemporaryDirectory() as tempdir:
+      sm_path = os.path.join(tempdir, "interpolate.savedmodel")
+      tf_integration.save_torch_module_as_tf_saved_model(pt_model, arg, sm_path)
+
+      # Reload SM and compare results with PT results
+      loaded_model = tf.saved_model.load(sm_path)
+      pt_res = pt_model(*arg)
+      tf_res = torch.tensor(loaded_model.f(*arg)[0].numpy())
+      self.assertTrue(torch.allclose(pt_res, tf_res, atol=1e-4))
+
+
+
+if __name__ == "__main__":
+  test_base.main()

experimental/torch_xla2/torch_xla2/tf_integration.py

@@ -9,16 +9,16 @@
 
 
 def exported_program_to_tf_function(ep, enable_xla=True):
-  jax_program = export.exported_program_to_jax_program(ep)
-
-  example_inputs = jax_program.flatten_inputs(*jax_program.example_inputs)
+  weights, jax_program = export.exported_program_to_jax(ep)
+  wrapped = lambda *args: jax_program(weights, (args,))
+  avals = export.extract_avals(ep)
   input_signature = [
       tf.TensorSpec(shape=t.shape, dtype=t.dtype, name=f"args_{i}")
-      for i, t in enumerate(example_inputs)
+      for i, t in enumerate(avals)
   ]
   tf_f = tf.function(
       jax2tf.convert(
-          jax_program.flatten_callable,
+          wrapped,
           with_gradient=False,
           enable_xla=enable_xla,
       ),