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,11 +2305,19 @@ 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());

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
@@ -68,11 +69,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;
   }
 }