Fix persistent compilation cache not working on TPU

torch_xla/csrc/runtime/pjrt_computation_client.cpp

@@ -697,17 +697,35 @@ std::string PjRtComputationClient::SerializeComputation(
 
 ComputationClient::ComputationPtr PjRtComputationClient::DeserializeComputation(
     const std::string& serialized) {
-  absl::StatusOr<std::unique_ptr<xla::PjRtExecutable>> executable_or =
-      client_->DeserializeExecutable(serialized, std::nullopt);
-  if (!executable_or.ok()) {
-    TF_LOG(WARNING) << "Failed to deserialize executable: "
-                    << executable_or.status();
-    return nullptr;
+  std::unique_ptr<xla::PjRtLoadedExecutable> loaded_executable;
+
+  // First, try LoadSerializedExecutable which directly produces a loaded
+  // executable. This is the path implemented by PJRT C API plugins (e.g. TPU).
+  absl::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable>>
+      loaded_executable_or = client_->LoadSerializedExecutable(
+          serialized, std::nullopt, xla::LoadOptions());
+  if (loaded_executable_or.ok()) {
+    loaded_executable = std::move(loaded_executable_or.value());
+  } else {
+    // Fall back to the two-step DeserializeExecutable + Load path for backends
+    // that implement DeserializeExecutable instead of LoadSerializedExecutable.
+    absl::StatusOr<std::unique_ptr<xla::PjRtExecutable>> executable_or =
+        client_->DeserializeExecutable(serialized, std::nullopt);
+    if (!executable_or.ok()) {
+      TF_LOG(WARNING) << "Failed to deserialize executable: "
+                      << loaded_executable_or.status() << " ; "
+                      << executable_or.status();
+      return nullptr;
+    }
+    absl::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable>> load_or =
+        client_->Load(std::move(executable_or.value()), xla::LoadOptions());
+    if (!load_or.ok()) {
+      TF_LOG(WARNING) << "Failed to load deserialized executable: "
+                      << load_or.status();
+      return nullptr;
+    }
+    loaded_executable = std::move(load_or.value());
   }
-  std::unique_ptr<xla::PjRtExecutable> executable =
-      std::move(executable_or.value());
-  std::unique_ptr<xla::PjRtLoadedExecutable> loaded_executable =
-      client_->Load(std::move(executable), xla::LoadOptions()).value();
 
   auto hlo_modules = loaded_executable->GetHloModules();
   if (!hlo_modules.ok()) {