Merge pull request #33492 from gnecula:export_memories

PiperOrigin-RevId: 836271707

Author: Google-ML-Automation

jax/_src/export/serialization.fbs

@@ -82,10 +82,18 @@ enum DType: byte {
   key_unsafe_rbg = 29,
 }
 
+enum MemorySpace: byte {
+  Missing = 0,  // default if missing (pre 11/25/2025)
+  Device = 1,
+  Host = 2,
+  Any = 3,
+}
+
 table AbstractValue {
   kind: AbstractValueKind;
   shape: [string];  // Support shape polymorphism
   dtype: DType;
+  memory_space: MemorySpace;
 }
 
 enum ShardingKind: byte {
@@ -119,14 +127,15 @@ table Exported {
   /// Note that this field has different semantics and purpose from
   /// `mlir_module_serialization_version`, which encodes
   /// the calling convention of the `mlir_module_serialized`.
+  /// See comments in serialization.py for more details.
   serialization_version: uint16;
 
   function_name: string;
   in_tree: PyTreeDef;
   in_avals: [AbstractValue];
   out_tree: PyTreeDef;
   out_avals: [AbstractValue];
-  nr_devices: short;
+  nr_devices_short: short;  // Deprecated as of 11/25/2025
   in_shardings: [Sharding];
   out_shardings: [Sharding];
 
@@ -142,6 +151,7 @@ table Exported {
   uses_global_constants: bool;
 
   vjp: Exported;
+  nr_devices: uint32 = 0;  // Added 11/25/2025
 }
 
 root_type Exported;

jax/_src/export/serialization.py

@@ -50,7 +50,15 @@
 #   This version is backwards compatible with Version 2.
 # Version 4, April 7th, 2025, adds serialization for PRNGs key types.
 #   This version is backwards compatible with Version 2 and 3.
-_SERIALIZATION_VERSION = 2
+# Version 5, November 23rd, 2025, adds serialization for aval memory_space,
+#   upgrade num_devices to a 32 bit value.
+#   This version is backwards compatible with Version 2 to 4.
+# TODO(necula): we cannot really store the actual serialization_version
+# in the flatbuffer because prior to 11/25/2025 deserializers checked
+# if the version is 2 or 3. I have now removed that check, but for the
+# sake of old deserializers we can only store version 3. Starting
+# on January 2026 we can store the actual version.
+_SERIALIZATION_VERSION = 3
 
 def serialize(exp: _export.Exported, vjp_order: int = 0) -> bytearray:
   """Serializes an Exported.
@@ -157,10 +165,6 @@ def _serialize_array(
 
 def _deserialize_exported(exp: ser_flatbuf.Exported) -> _export.Exported:
   serialization_version = exp.SerializationVersion()
-  if serialization_version not in [2, 3]:
-    raise NotImplementedError(
-        f"deserialize unsupported version {serialization_version}"
-    )
 
   fun_name = exp.FunctionName().decode("utf-8")
   in_tree = tree_util.tree_structure(
@@ -177,7 +181,10 @@ def _deserialize_exported(exp: ser_flatbuf.Exported) -> _export.Exported:
   out_avals = _deserialize_tuple(
       exp.OutAvalsLength, exp.OutAvals, deser_aval
   )
-  nr_devices = exp.NrDevices()
+  # TODO(necula): remove the fallback to NrDevicesShort and mark
+  # the field "deprecated" once we abandon the old
+  # serialization format (6 months after 11/24/2025).
+  nr_devices = exp.NrDevices() or exp.NrDevicesShort()
   in_shardings = _deserialize_tuple(
       exp.InShardingsLength, exp.InShardings, _deserialize_sharding
   )
@@ -381,6 +388,14 @@ def register_dtype_kind(dtype: Any, kind: int):
   _dtype_kind_to_dtype[kind] = dtype
 
 
+_memory_space_to_enum = {
+    core.MemorySpace.Device: ser_flatbuf.MemorySpace.Device,
+    core.MemorySpace.Host: ser_flatbuf.MemorySpace.Host,
+    core.MemorySpace.Any: ser_flatbuf.MemorySpace.Any,
+}
+_memory_space_from_enum = {v: k for k, v in _memory_space_to_enum.items()}
+
+
 def _serialize_aval(
     builder: flatbuffers.Builder, aval: core.ShapedArray
 ) -> int:
@@ -395,6 +410,7 @@ def _serialize_aval(
   ser_flatbuf.AbstractValueAddKind(builder, aval_kind)
   ser_flatbuf.AbstractValueAddShape(builder, shape_vector_offset)
   ser_flatbuf.AbstractValueAddDtype(builder, _dtype_to_dtype_kind[aval.dtype])
+  ser_flatbuf.AbstractValueAddMemorySpace(builder, _memory_space_to_enum[aval.memory_space])
   return ser_flatbuf.AbstractValueEnd(builder)
 
 
@@ -409,7 +425,8 @@ def _deserialize_aval(aval: ser_flatbuf.AbstractValue,
         ),
         scope=scope
     )
-    return core.ShapedArray(shape, dtype)
+    mem_space = aval.MemorySpace() or ser_flatbuf.MemorySpace.Device
+    return core.ShapedArray(shape, dtype, memory_space=_memory_space_from_enum[mem_space])
   else:
     assert False, aval_kind
 

jax/_src/export/serialization_generated.py

@@ -21,7 +21,7 @@
 from flatbuffers.compat import import_numpy
 np = import_numpy()
 
-class PyTreeDefKind:
+class PyTreeDefKind(object):
     leaf = 0
     none = 1
     tuple = 2
@@ -30,12 +30,12 @@ class PyTreeDefKind:
     custom = 5
 
 
-class AbstractValueKind:
+class AbstractValueKind(object):
     shapedArray = 0
     abstractToken = 1
 
 
-class DType:
+class DType(object):
     bool = 0
     i8 = 1
     i16 = 2
@@ -68,18 +68,25 @@ class DType:
     key_unsafe_rbg = 29
 
 
-class ShardingKind:
+class MemorySpace(object):
+    Missing = 0
+    Device = 1
+    Host = 2
+    Any = 3
+
+
+class ShardingKind(object):
     unspecified = 0
     hlo_sharding = 1
 
 
-class DisabledSafetyCheckKind:
+class DisabledSafetyCheckKind(object):
     platform = 0
     custom_call = 1
     shape_assertions = 2
 
 
-class PyTreeDef:
+class PyTreeDef(object):
     __slots__ = ['_tab']
 
     @classmethod
@@ -214,7 +221,7 @@ def PyTreeDefEnd(builder):
 
 
 
-class AbstractValue:
+class AbstractValue(object):
     __slots__ = ['_tab']
 
     @classmethod
@@ -266,8 +273,15 @@ def Dtype(self):
             return self._tab.Get(flatbuffers.number_types.Int8Flags, o + self._tab.Pos)
         return 0
 
+    # AbstractValue
+    def MemorySpace(self):
+        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(10))
+        if o != 0:
+            return self._tab.Get(flatbuffers.number_types.Int8Flags, o + self._tab.Pos)
+        return 0
+
 def AbstractValueStart(builder):
-    builder.StartObject(3)
+    builder.StartObject(4)
 
 def AbstractValueAddKind(builder, kind):
     builder.PrependInt8Slot(0, kind, 0)
@@ -281,12 +295,15 @@ def AbstractValueStartShapeVector(builder, numElems):
 def AbstractValueAddDtype(builder, dtype):
     builder.PrependInt8Slot(2, dtype, 0)
 
+def AbstractValueAddMemorySpace(builder, memorySpace):
+    builder.PrependInt8Slot(3, memorySpace, 0)
+
 def AbstractValueEnd(builder):
     return builder.EndObject()
 
 
 
-class Sharding:
+class Sharding(object):
     __slots__ = ['_tab']
 
     @classmethod
@@ -355,7 +372,7 @@ def ShardingEnd(builder):
 
 
 
-class Effect:
+class Effect(object):
     __slots__ = ['_tab']
 
     @classmethod
@@ -391,7 +408,7 @@ def EffectEnd(builder):
 
 
 
-class DisabledSafetyCheck:
+class DisabledSafetyCheck(object):
     __slots__ = ['_tab']
 
     @classmethod
@@ -437,7 +454,7 @@ def DisabledSafetyCheckEnd(builder):
 
 
 
-class Exported:
+class Exported(object):
     __slots__ = ['_tab']
 
     @classmethod
@@ -460,6 +477,7 @@ def Init(self, buf, pos):
     # Note that this field has different semantics and purpose from
     # `mlir_module_serialization_version`, which encodes
     # the calling convention of the `mlir_module_serialized`.
+    # See comments in serialization.py for more details.
     # Exported
     def SerializationVersion(self):
         o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
@@ -543,7 +561,7 @@ def OutAvalsIsNone(self):
         return o == 0
 
     # Exported
-    def NrDevices(self):
+    def NrDevicesShort(self):
         o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(16))
         if o != 0:
             return self._tab.Get(flatbuffers.number_types.Int16Flags, o + self._tab.Pos)
@@ -767,8 +785,15 @@ def Vjp(self):
             return obj
         return None
 
+    # Exported
+    def NrDevices(self):
+        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(40))
+        if o != 0:
+            return self._tab.Get(flatbuffers.number_types.Uint32Flags, o + self._tab.Pos)
+        return 0
+
 def ExportedStart(builder):
-    builder.StartObject(18)
+    builder.StartObject(19)
 
 def ExportedAddSerializationVersion(builder, serializationVersion):
     builder.PrependUint16Slot(0, serializationVersion, 0)
@@ -794,8 +819,8 @@ def ExportedAddOutAvals(builder, outAvals):
 def ExportedStartOutAvalsVector(builder, numElems):
     return builder.StartVector(4, numElems, 4)
 
-def ExportedAddNrDevices(builder, nrDevices):
-    builder.PrependInt16Slot(6, nrDevices, 0)
+def ExportedAddNrDevicesShort(builder, nrDevicesShort):
+    builder.PrependInt16Slot(6, nrDevicesShort, 0)
 
 def ExportedAddInShardings(builder, inShardings):
     builder.PrependUOffsetTRelativeSlot(7, flatbuffers.number_types.UOffsetTFlags.py_type(inShardings), 0)
@@ -854,5 +879,8 @@ def ExportedAddUsesGlobalConstants(builder, usesGlobalConstants):
 def ExportedAddVjp(builder, vjp):
     builder.PrependUOffsetTRelativeSlot(17, flatbuffers.number_types.UOffsetTFlags.py_type(vjp), 0)
 
+def ExportedAddNrDevices(builder, nrDevices):
+    builder.PrependUint32Slot(18, nrDevices, 0)
+
 def ExportedEnd(builder):
     return builder.EndObject()

tests/export_test.py

@@ -1405,6 +1405,18 @@ def f_with_sharding(x):
         f"context with {jax.local_device_count()} devices"):
       exp.call(b)
 
+  def test_memory_space(self):
+    shd = jax.sharding.SingleDeviceSharding(
+        jax.devices()[0], memory_kind="pinned_host")
+    a = jax.device_put(1, shd)
+    f = jax.jit(lambda x: x)
+
+    exported = get_exported(f, platforms=("tpu", "cuda"))(a)
+    self.assertEqual(exported.in_avals[0].memory_space, core.MemorySpace.Host)
+    if jtu.device_under_test in ("tpu", "gpu"):
+      b = exported.call(a)
+      self.assertEqual(b.sharding, a.sharding)
+
   @jtu.parameterized_filterable(
     kwargs=[
       dict(testcase_name=f"_poly={poly}", poly=poly)

tests/memories_test.py

@@ -743,23 +743,6 @@ def test_host_to_device_transfer(self):
       self.assertEqual(d.sharding.memory_kind, 'device')
       self.assertArraysEqual(d, orig)
 
-  def test_memory_space_propagated_identity_jit(self):
-    shd = jax.sharding.SingleDeviceSharding(
-        jax.devices()[0], memory_kind='pinned_host')
-    a = jax.device_put(1, shd)
-
-    f = jax.jit(lambda x: x, out_shardings=shd)
-    b = f(a)
-    self.assertEqual(b.sharding, a.sharding)
-
-    f = jax.jit(lambda x: x)
-    b = f(a)
-    self.assertEqual(b.sharding, a.sharding)
-
-    exported = jax.export.export(f)(a)
-    b = exported.call(a)
-    self.assertEqual(b.sharding, a.sharding)
-
 
 class ComputeOffload(jtu.BufferDonationTestCase):