test inplace update together with xm.save

test/test_inplace_update.py

@@ -1,10 +1,11 @@
-import os
+import io
 import sys
 import unittest
 
 import torch
 import torch_xla
 import torch_xla.core.xla_model as xm
+from test_utils import temporary_env
 
 
 class InplaceUpdateTest(unittest.TestCase):
@@ -49,6 +50,26 @@ def test_non_aten_op_after_partial_update(self):
     xm.mark_step()
     self.assertTrue(torch.all(torch.eq(y, expected)))
 
+  def test_xm_save(self):
+    with temporary_env(
+        XLA_DISABLE_FUNCTIONALIZATION="0", XLA_ENABLE_PARAM_ALIASING="0"):
+      xla_device = xm.xla_device()
+      t1 = torch.tensor([1], device=xla_device)
+      t2 = t1.detach()
+      xm.mark_step()
+
+      t2.add_(t2)
+      xm.mark_step()
+
+      # mark_step() causes t1 and t2 to be out of sync on the XLA side.
+
+      fobj = io.BytesIO()
+      xm.save({'t1': t1}, fobj)
+      fobj.seek(0)
+      saved = torch.load(fobj)
+
+      self.assertEqual(t1.item(), saved['t1'].item())
+
 
 if __name__ == '__main__':
   test = unittest.main()

test/test_utils.py

@@ -1,4 +1,5 @@
 import collections
+from contextlib import contextmanager
 import itertools
 import math
 import os
@@ -390,3 +391,32 @@ def runAtenTest(self, tensors, fn, device=None, rel_err=1e-2, abs_err=1e-5):
     results = xu.as_list(fn(*tensors))
     xla_results = xu.as_list(fn(*xla_tensors))
     self.compareResults(results, xla_results, rel_err=rel_err, abs_err=abs_err)
+
+
+@contextmanager
+def temporary_env(**kwargs):
+  """
+    Temporarily set environment variables within the context.
+    
+    Args:
+        **kwargs: Key-value pairs representing environment variables to set.
+                  For example: temporary_env(PATH='/new/path', DEBUG='1')
+    """
+  original_env = {}
+
+  # Store original values and set new ones
+  for key, value in kwargs.items():
+    original_env[key] = os.environ.get(key, None)
+    os.environ[key] = value
+
+  try:
+    yield
+  finally:
+    # Restore original environment variables
+    for key, old_value in original_env.items():
+      if old_value is None:
+        # The variable was not originally set
+        del os.environ[key]
+      else:
+        # Restore the original value
+        os.environ[key] = old_value