[PJRT C API] Check whether the PJRT_Api* for the device type already …

…exists before calling dlopen and dlsym.

PiperOrigin-RevId: 531295150

jax/_src/xla_bridge.py

@@ -328,7 +328,12 @@ def factory():
         library_path = path
         options = None
 
-      xla_client.load_pjrt_plugin_dynamically(name, library_path)
+      if xla_extension_version >= 152:
+        # Plugin may already be statically linked in some configurations.
+        if not xla_client.pjrt_plugin_loaded(name):
+          xla_client.load_pjrt_plugin_dynamically(name, library_path)
+      else:
+        xla_client.load_pjrt_plugin_dynamically(name, library_path)
       return xla_client.make_c_api_client(name, options)
 
     return factory

tests/xla_bridge_test.py

@@ -98,7 +98,13 @@ def test_register_plugin(self):
       with mock.patch.object(
           xc, "load_pjrt_plugin_dynamically", autospec=True
       ) as mock_load_plugin:
-        client_factory()
+        if xc._version >= 152:
+          with mock.patch.object(
+              xc, "pjrt_plugin_loaded", autospec=True
+          ) as mock_plugin_loaded:
+            client_factory()
+        else:
+          client_factory()
 
     self.assertRegex(
         log_output[1][0],
@@ -108,7 +114,10 @@ def test_register_plugin(self):
     self.assertIn("name1", xb._backend_factories)
     self.assertIn("name2", xb._backend_factories)
     self.assertEqual(priotiy, 400)
-    mock_load_plugin.assert_called_once_with("name1", "path1")
+    if xc._version >= 152:
+      mock_plugin_loaded.assert_called_once_with("name1")
+    else:
+      mock_load_plugin.assert_called_once_with("name1", "path1")
     if xc._version >= 134:
       mock_make.assert_called_once_with("name1", None)
     else:
@@ -126,13 +135,22 @@ def test_register_plugin_with_config(self):
       with mock.patch.object(
           xc, "load_pjrt_plugin_dynamically", autospec=True
       ) as mock_load_plugin:
-        client_factory()
+        if xc._version >= 152:
+          with mock.patch.object(
+              xc, "pjrt_plugin_loaded", autospec=True
+          ) as mock_plugin_loaded:
+            client_factory()
+        else:
+          client_factory()
 
     self.assertIn("name1", xb._backend_factories)
     self.assertEqual(priority, 400)
-    mock_load_plugin.assert_called_once_with(
-        "name1", "/path/pjrt_plugin_name1.so"
-    )
+    if xc._version >= 152:
+      mock_plugin_loaded.assert_called_once_with("name1")
+    else:
+      mock_load_plugin.assert_called_once_with(
+          "name1", "/path/pjrt_plugin_name1.so"
+      )
     mock_make.assert_called_once_with(
         "name1",
         {
@@ -147,6 +165,7 @@ def test_register_plugin_with_config(self):
 class GetBackendTest(jtu.JaxTestCase):
 
   class _DummyBackend:
+
     def __init__(self, platform, device_count):
       self.platform = platform
       self._device_count = device_count