Merge pull request #33848 from gnecula:export_memory_space

PiperOrigin-RevId: 842754890

Author: Google-ML-Automation

jax/_src/internal_test_util/export_back_compat_test_data/export_with_memory_space.py

@@ -0,0 +1,23 @@
+# Copyright 2025 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# ruff: noqa
+
+# Pasted from the test output (see export_serialization_back_compat_test.py module docstring)
+serializations = [
+    dict(
+        serialization_version=5,
+        exported_serialized=bytearray(b"0\x00\x00\x00\x00\x00*\x00L\x00J\x00D\x00@\x00<\x008\x004\x00.\x00(\x00$\x00 \x00\x1c\x00\x18\x00\x14\x00\x10\x00\x0e\x00\x08\x00\x07\x00\x00\x000\x00*\x00\x00\x00\x00\x00\x00\x01D\x00\x00\x00\x00\x00\n\x00D\x00\x00\x00\x84\x02\x00\x00\x84\x02\x00\x00\x84\x02\x00\x00X\x02\x00\x00\x80\x02\x00\x00\x88\x02\x00\x00\x00\x00\x02\x00\x02\x00\x00\x00\xa0\x02\x00\x00\xcc\x02\x00\x00\xcc\x02\x00\x00\x04\x03\x00\x00X\x03\x00\x00\x00\x00\x03\x00\x01\x00\x00\x00\x00\x00\x00\x00 \x02\x00\x00ML\xefR\rStableHLO_v1.13.0\x00\x01\x1f\x07\x01\x05\t\t\x01\x03\x0f\x03\x03\x13\x05\x05\x17\x1b\x03kE\x0f\x01\x1b\x07\x0b#\x0b\x0f\x0b\x0f\x0b\x0f\x0b\x0b\x0f\x0b\x03\r\x13\x0f\x1b\x17\x0f\x13\x05\x1f\x0b\x0b\x13\x13\x0b\x0b\x1b\x0b\x0b\x0f\x1b\x0b\x0b\x0b\x0b\x01\x05\x0f\x0b\x05\x0b\x17\x0f\x1b\x07\x07\x02\xf9\x1f\x05\t\x03\x07\x07\t\x0b\r\x0f\x11\x05\x0f\x11\x03\x01\x05\x11\x11\x01\t\x05\x13\x11\x01\x05\x05\x15\t\x03\x1d\x19\x01\x05\x17\x05\x03\x1d\x01\x03\x17\t\r\x15\x05!%\x01\x0b\x03#\x01\x01\t\x17\x01\x0b\x01\x01\x01\x1d\x19\x1d\x1b\x03\x05-3\r\x03/1\x1d\x1d\x1d\x1f\r\x05\')5\x1f\x1d!#\t\x03\x03;\r\x05=?\')\x1d#\x1d%\x1d\'\x1d)\x01\x02\x02\x01\t)\x05\t\r\r)\x01\x0b\x11\x05\x07\x05\x03\x05\x1b\t\x04I\x05\x01Q\x01\x05\x01\x07\x047\x03\x01\t\x03@\x01\x03\x05P\x01\x05\x07\x04\x1b\x03\x05\x07\x05\r\x0b\x17\x00\x07\x04\x01\x03\x03\x06\x03\x01\x05\x01\x00\x1a\x04+\x0f\x0b\x0f!\x1b!)\x19#\x05\x19%)9\x15\x11\x0b\x0f\x0b\t\x11builtin\x00sdy\x00vhlo\x00module\x00mesh\x00func_v1\x00return_v1\x00jax.uses_shape_polymorphism\x00mhlo.num_partitions\x00mhlo.num_replicas\x00jit__lambda\x00x\x00mhlo.memory_kind\x00pinned_host\x00jax.global_constant\x00_platform_index\x00sdy.sharding\x00jax.result_info\x00result\x00main\x00public\x00\x08\x1b\x07\x05\'\x01\x05\x1b\x03\x0b+79AC\x02\x00\x00\x00\x14\x00\x00\x00\x04\x00\x00\x00\x04\x00\x00\x00cuda\x00\x00\x00\x00\x03\x00\x00\x00tpu\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x04\x00\x00\x00\x18\xff\xff\xff\x01\x00\x00\x00\x04\x00\x00\x00@\xff\xff\xff\x08\x00\x00\x00\x00\x00\x00\x01\x0c\x00\x00\x00\x08\x03\x1a\x02\x02\x01J\x01\x02R\x01\x00\x01\x00\x00\x00\x04\x00\x00\x00\xcc\xff\xff\xff\x00\x00\x02\n\x04\x00\x00\x00\x02\x00\x00\x00\x10\x00\x00\x00\x04\x00\x00\x00\x01\x00\x00\x003\x00\x00\x00\x01\x00\x00\x002\x00\x00\x00p\xff\xff\xff\x01\x00\x00\x00\x10\x00\x00\x00\x0c\x00\x0c\x00\x00\x00\x08\x00\x07\x00\x06\x00\x0c\x00\x00\x00\x00\x00\x02\n\x04\x00\x00\x00\x02\x00\x00\x00\x10\x00\x00\x00\x04\x00\x00\x00\x01\x00\x00\x003\x00\x00\x00\x01\x00\x00\x002\x00\x00\x00\xcc\xff\xff\xff\x08\x00\x00\x00\x00\x00\x00\x02\x02\x00\x00\x00,\x00\x00\x00\x10\x00\x00\x00\x00\x00\n\x00\x0c\x00\x0b\x00\x00\x00\x04\x00\n\x00\x00\x00\x08\x00\x00\x00\x00\x00\x00\x04\x00\x00\x00\x00\x08\x00\x0c\x00\x0b\x00\x04\x00\x08\x00\x00\x00\x08\x00\x00\x00\x00\x00\x00\x02\x01\x00\x00\x00\x08\x00\x00\x00\x04\x00\x04\x00\x04\x00\x00\x00\x08\x00\x00\x00<lambda>\x00\x00\x00\x00"),
+    ),
+]

tests/export_serialization_back_compat_test.py

@@ -23,8 +23,8 @@
   * Create a new test method, with a function to be serialized that exercises
     the feature you want to test, and a call to self.export_and_serialize.
     You can follow the model of the tests below, which are parameterized by
-    the test data. Use `None` for the test data to signal that you want to
-    use a fresh serialization.
+    the testdata. Use only `None` for the testdata parameter to signal that
+    you want to use a current serialization and not a saved one.
   * Run the test. This will save the serialized data in
     TEST_UNDECLARED_OUTPUTS_DIR (or "/tmp/back_compat_testdata" if not set).
   * Copy the test data defined in the output file, to the file
@@ -55,13 +55,15 @@
 
 import jax
 from jax._src import config
+from jax._src import core
 from jax._src.export import _export
 from jax._src.export.serialization import _SERIALIZATION_VERSION
 from jax.sharding import PartitionSpec as P
 from jax._src import test_util as jtu
 
 from jax._src.internal_test_util.export_back_compat_test_data import export_with_specified_sharding
 from jax._src.internal_test_util.export_back_compat_test_data import export_with_unspecified_sharding
+from jax._src.internal_test_util.export_back_compat_test_data import export_with_memory_space
 
 config.parse_flags_with_absl()
 jtu.request_cpu_devices(8)
@@ -75,14 +77,15 @@ def setUp(self):
 
   def export_and_serialize(self, fun, *args,
                            vjp_order=0,
+                           platforms=None,
                            **kwargs) -> bytearray:
     """Export and serialize a function.
 
     The test data is saved in TEST_UNDECLARED_OUTPUTS_DIR (or
     "/tmp/back_compat_testdata" if not set) and should be copied as explained
     in the module docstring.
     """
-    exp = _export.export(fun)(*args, **kwargs)
+    exp = _export.export(fun, platforms=platforms)(*args, **kwargs)
     serialized = exp.serialize(vjp_order=vjp_order)
     updated_testdata = f"""
     # Paste to the test data file (see export_serialization_back_compat_test.py module docstring)
@@ -98,7 +101,8 @@ def export_and_serialize(self, fun, *args,
                            "/tmp/back_compat_testdata")
     if not os.path.exists(output_dir):
       os.makedirs(output_dir)
-    output_file = os.path.join(output_dir, f"export_{self._testMethodName}.py")
+    output_file_basename = f"export_{self._testMethodName.replace('test_', '')}.py"
+    output_file = os.path.join(output_dir, output_file_basename)
     logging.info("Writing the updated serialized Exported at %s", output_file)
     with open(output_file, "w") as f:
       f.write(updated_testdata)
@@ -163,5 +167,38 @@ def f(b):
       self.assertEqual(out.addressable_shards[1].index, (slice(8, 16), slice(None)))
 
 
+  @jtu.parameterized_filterable(
+    kwargs=[
+      dict(testdata=testdata,
+      testcase_name=("current" if testdata is None
+                     else f"v{testdata['serialization_version']}"))
+      for testdata in [None, *export_with_memory_space.serializations]
+    ]
+  )
+  def test_with_memory_space(self, testdata: dict[str, Any] | None):
+    # This test is based on export_test.py::test_memory_space_from_arg
+    mesh = jtu.create_mesh((2,), "x")
+    with jax.set_mesh(mesh):
+      shd = jax.sharding.NamedSharding(mesh, P("x", None),
+                                       memory_kind="pinned_host")
+      a = jax.device_put(np.ones((2, 3), dtype=np.float32), shd)
+      f = jax.jit(lambda x: x)
+
+      if testdata is None:
+        serialized = self.export_and_serialize(
+            f, a, platforms=("tpu", "cuda"))
+      else:
+        serialized = testdata["exported_serialized"]
+
+    exported = _export.deserialize(serialized)
+    self.assertEqual(exported.in_avals[0].memory_space, core.MemorySpace.Host)
+    self.assertEqual(exported.out_avals[0].memory_space, core.MemorySpace.Host)
+
+    if jtu.device_under_test() in ("tpu", "gpu"):
+      b = exported.call(a)
+      self.assertEqual(b.aval.memory_space, core.MemorySpace.Host)
+      self.assertEqual(b.sharding, a.sharding)
+
+
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())