Replace std::shared_ptr<xla::ifrt::LoadedExecutable> with xla::ifrt::LoadedExecutableRef

PiperOrigin-RevId: 755581247

Author: junwhanahn

jaxlib/py_compile_only_client.cc

@@ -70,7 +70,7 @@ class CompileOnlyPyClient : public PyClient {
     return client;
   }
 
-  absl::StatusOr<std::shared_ptr<ifrt::Executable>> CompileUnloaded(
+  absl::StatusOr<ifrt::ExecutableRef> CompileUnloaded(
       absl::string_view mlir_module, ifrt::DeviceListRef executable_devices,
       CompileOptions options, std::vector<nb::capsule> host_callbacks) {
     if (!host_callbacks.empty()) {
@@ -102,7 +102,7 @@ class CompileOnlyPyClient : public PyClient {
                                     *ifrt_client->topology().description()));
     TF_ASSIGN_OR_RETURN(auto ifrt_executable,
                         ifrt::PjRtExecutable::Create(std::move(executable)));
-    return std::shared_ptr<ifrt::Executable>(std::move(ifrt_executable));
+    return ifrt::ExecutableRef(std::move(ifrt_executable));
   }
 
  private:

jaxlib/py_executable.cc

@@ -84,7 +84,7 @@ absl::Status PyShardedToken::Await() {
 
 PyLoadedExecutable::PyLoadedExecutable(
     nb_class_ptr<PyClient> client,
-    std::shared_ptr<ifrt::LoadedExecutable> ifrt_loaded_executable,
+    ifrt::LoadedExecutableRef ifrt_loaded_executable,
     std::optional<nb_traceback> traceback,
     std::optional<std::string> fingerprint)
     : client_(std::move(client)),

jaxlib/py_executable.h

@@ -131,19 +131,18 @@ using ExecuteShardedArg = std::variant<PyArray, std::vector<PyArray>>;
 // b) to add Python-specific functionality.
 class PyLoadedExecutable {
  public:
-  PyLoadedExecutable(
-      nb_class_ptr<PyClient> client,
-      std::shared_ptr<ifrt::LoadedExecutable> ifrt_loaded_executable,
-      std::optional<nb_traceback> traceback,
-      std::optional<std::string> fingerprint);
+  PyLoadedExecutable(nb_class_ptr<PyClient> client,
+                     ifrt::LoadedExecutableRef ifrt_loaded_executable,
+                     std::optional<nb_traceback> traceback,
+                     std::optional<std::string> fingerprint);
   ~PyLoadedExecutable();
 
   nb_class_ptr<PyClient> client() const { return client_; }
   ifrt::LoadedExecutable* ifrt_loaded_executable() const {
     return ifrt_loaded_executable_.get();
   }
 
-  std::shared_ptr<ifrt::LoadedExecutable> shared_ifrt_loaded_executable() {
+  ifrt::LoadedExecutableRef shared_ifrt_loaded_executable() {
     return ifrt_loaded_executable_;
   }
 
@@ -226,7 +225,7 @@ class PyLoadedExecutable {
   friend class PyClient;
 
   nb_class_ptr<PyClient> client_;
-  std::shared_ptr<ifrt::LoadedExecutable> ifrt_loaded_executable_;
+  ifrt::LoadedExecutableRef ifrt_loaded_executable_;
   std::optional<nb_traceback> traceback_;
 
   // Identical executables (i.e. representing the same program) will have the