Add platform name to xla::ifrt::Device

PiperOrigin-RevId: 844067251

Author: ICGog

jaxlib/py_array.cc

@@ -47,6 +47,7 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
+#include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "llvm/Support/Casting.h"
 #include "nanobind/nanobind.h"
@@ -2304,13 +2305,21 @@ absl::Status PyArray::Register(nb::module_& m) {
       nb::is_method());
   type.attr("platform") = nb::cpp_function(
       [](PyArray self) {
-        if (self.ifrt_array()->client()->platform_name() == "cuda" ||
-            self.ifrt_array()->client()->platform_name() == "rocm") {
+#if JAX_IFRT_VERSION_NUMBER >= 43
+        const xla::ifrt::DeviceListRef& devices =
+            self.ifrt_array()->sharding().devices();
+        absl::string_view platform_name =
+            devices->devices().front()->PlatformName();
+#else
+        absl::string_view platform_name =
+            self.ifrt_array()->client()->platform_name();
+#endif
+        if (platform_name == "cuda" || platform_name == "rocm") {
           return std::string_view("gpu");
         } else {
-          return self.ifrt_array()->client()->platform_name();
+          return platform_name;
         }
-      },
+    },
       nb::is_method());
   type.attr("is_ready") = nb::cpp_function(
       [](PyArray self) { return xla::ValueOrThrow(self.IsReady()); },

jaxlib/py_device.cc

@@ -28,6 +28,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_join.h"
+#include "absl/strings/string_view.h"
 #include "llvm/Support/Casting.h"
 #include "nanobind/nanobind.h"
 #include "nanobind/stl/optional.h"  // IWYU pragma: keep
@@ -45,6 +46,7 @@ limitations under the License.
 #include "xla/python/nb_helpers.h"
 #include "xla/python/pjrt_ifrt/pjrt_client.h"
 #include "xla/python/pjrt_ifrt/pjrt_device.h"
+#include "xla/python/version.h"
 #include "xla/tsl/framework/allocator.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
@@ -68,11 +70,15 @@ std::string_view PyDevice::platform() const {
   // but we haven't yet updated JAX clients that
   // expect "gpu". Migrate users and remove this
   // code.
-  if (client_->platform_name() == "cuda" ||
-      client_->platform_name() == "rocm") {
+#if JAX_IFRT_VERSION_NUMBER >= 43
+  absl::string_view platform_name = device_->PlatformName();
+#else
+  absl::string_view platform_name = client_->platform_name();
+#endif
+  if (platform_name == "cuda" || platform_name == "rocm") {
     return std::string_view("gpu");
   } else {
-    return client_->platform_name();
+    return platform_name;
   }
 }