Add LoadedExecutable.serialize() to match Executable.serialize().

PiperOrigin-RevId: 841880441

Author: pschuh

jax/_src/compilation_cache.py

@@ -261,7 +261,10 @@ def put_executable_and_time(
                " since cache is disabled/not initialized", cache_key)
     return
 
-  serialized_executable = backend.serialize_executable(executable)
+  if hasattr(executable, "serialize") or xla_client._version >= 389:
+    serialized_executable = executable.serialize()
+  else:
+    serialized_executable = backend.serialize_executable(executable)
   executable_and_time = combine_executable_and_time(
       serialized_executable, compile_time)
   executable_and_time = compress_executable(executable_and_time)

jaxlib/_jax/__init__.pyi

@@ -1199,6 +1199,7 @@ class LoadedExecutable:
   def client(self) -> Client: ...
   def local_devices(self) -> list[Device]: ...
   def get_hlo_text(self) -> str: ...
+  def serialize(self) -> bytes: ...
   def size_of_generated_code_in_bytes(self) -> int: ...
   def get_compiled_memory_stats(self) -> CompiledMemoryStats: ...
   def execute_sharded(

jaxlib/py_executable.cc

@@ -545,6 +545,12 @@ void PyLoadedExecutable::Register(nb::module_& m) {
       .def("get_hlo_text",
            xla::ValueOrThrowWrapper(
                &PyLoadedExecutable::GetHumanReadableProgramText))
+      .def("serialize",
+           [](const PyLoadedExecutable& exec) -> nb::bytes {
+             std::string serialized =
+                 xla::ValueOrThrow(exec.ifrt_loaded_executable()->Serialize());
+             return nb::bytes(serialized.data(), serialized.size());
+           })
       .def("size_of_generated_code_in_bytes",
            &PyLoadedExecutable::SizeOfGeneratedCodeInBytes)
       .def(

jaxlib/xla_client.py

@@ -47,7 +47,7 @@
 # Please suffix the version number with a brief description of your change
 # in a comment. The goal here is to force a merge conflict if two changes
 # attempt to grab the same version number.
-_version = 388  # Add ArrayMeta
+_version = 389  # LoadedExecutable.serialize
 
 # An internal increasing version number for protecting jaxlib code against
 # ifrt changes.