[jax] Expose Device.local_hardware_id

This is to help us implement the new __dlpack_device__ interface
(https://dmlc.github.io/dlpack/latest/python_spec.html). I'm also ok
making this a private method, but figured it makes sense as a public
Device method.

PiperOrigin-RevId: 557286157

Author: skye

tensorflow/compiler/xla/python/xla.cc

@@ -201,6 +201,18 @@ PYBIND11_MODULE(xla_extension, m) {
                              [](const ClientAndPtr<PjRtDevice>& device) {
                                return device.client();
                              })
+      .def_property_readonly(
+          "local_hardware_id",
+          [](const ClientAndPtr<PjRtDevice>& device) -> std::optional<int> {
+            int local_hardware_id = device->local_hardware_id();
+            if (local_hardware_id == -1) {
+              return std::nullopt;
+            }
+            return local_hardware_id;
+          },
+          "Opaque hardware ID, e.g., the CUDA device number. In general, not "
+          "guaranteed to be dense, and not guaranteed to be defined on all "
+          "platforms.")
       .def("__str__", &PjRtDevice::DebugString)
       .def("__repr__", &PjRtDevice::ToString)
       .def("transfer_to_infeed",

tensorflow/compiler/xla/python/xla_client.py

@@ -44,7 +44,7 @@
 
 # Just an internal arbitrary increasing number to help with backward-compatible
 # changes. In JAX, reference this via jax._src.lib.xla_extension_version.
-_version = 183
+_version = 184
 
 # Version number for MLIR:Python components.
 mlir_api_version = 54

tensorflow/compiler/xla/python/xla_client_test.py

@@ -2215,6 +2215,12 @@ def testPlatform(self):
       for device in self.backend.local_devices():
         self.assertEqual(device.platform, self.backend.platform)
 
+    def testLocalHardwareId(self):
+      for device in self.backend.devices():
+        local_hardware_id = device.local_hardware_id
+        if local_hardware_id is not None:
+          self.assertGreaterEqual(local_hardware_id, 0)
+
     @unittest.skipIf(pathways, "not implemented")
     def testMemoryStats(self):
       for device in self.backend.local_devices():

tensorflow/compiler/xla/python/xla_extension/__init__.pyi

@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
+
 import enum
 import inspect
 import types
@@ -350,6 +352,7 @@ class Device:
   platform: str
   device_kind: str
   client: Client
+  local_hardware_id: int | None
   def __repr__(self) -> str: ...
   def __str__(self) -> str: ...
   def transfer_to_infeed(self, literal: _LiteralSlice): ...