[Mosaic GPU] Wire up the plugin that contains the MGPU custom call into Bazel tests

PiperOrigin-RevId: 751445068

Author: apaszke

jax_plugins/cuda/BUILD.bazel

@@ -12,15 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-licenses(["notice"])
-
 load(
-  "//jaxlib:jax.bzl",
-  "if_windows",
-  "py_library_providing_imports_info",
-  "pytype_library",
+    "//jaxlib:jax.bzl",
+    "if_windows",
+    "py_library_providing_imports_info",
+    "pytype_library",
 )
 
+licenses(["notice"])
+
 package(
     default_applicable_licenses = [],
     default_visibility = ["//:__subpackages__"],
@@ -41,7 +41,7 @@ py_library_providing_imports_info(
     ],
     data = if_windows(
         ["@xla//xla/pjrt/c/pjrt_c_api_gpu_plugin.pyd"],
-        ["@xla//xla/pjrt/c:pjrt_c_api_gpu_plugin.so"],
+        ["//jaxlib/tools:pjrt_c_api_gpu_plugin.so"],
     ),
     lib_rule = pytype_library,
 )

jax_plugins/cuda/__init__.py

@@ -51,7 +51,7 @@ def _get_library_path():
     runfiles_dir = os.getenv('RUNFILES_DIR', None)
     if runfiles_dir:
       local_path = os.path.join(
-          runfiles_dir, 'xla/xla/pjrt/c/pjrt_c_api_gpu_plugin.so'
+          runfiles_dir, '__main__/jaxlib/tools/pjrt_c_api_gpu_plugin.so'
       )
 
   if os.path.exists(local_path):