Add a jax network transfer benchmark script

PiperOrigin-RevId: 841983057

Author: Google-ML-Automation

jax/BUILD

@@ -459,7 +459,7 @@ alias(
 alias(
     name = "experimental_transfer",
     actual = "//jax/experimental:transfer",
-    visibility = [":internal"],
+    visibility = ["//jax/experimental:experimental_transfer_users"],
 )
 
 # Aliases of example_library targets.

jax/experimental/BUILD

@@ -15,6 +15,7 @@
 load(
     "//jaxlib:jax.bzl",
     "buffer_callback_internal_users",
+    "experimental_transfer_users",
     "if_cuda_is_configured",
     "jax_visibility",
     "mosaic_gpu_internal_users",
@@ -41,6 +42,12 @@ package_group(
     packages = buffer_callback_internal_users,
 )
 
+package_group(
+    name = "experimental_transfer_users",
+    includes = ["//jax:internal"],
+    packages = experimental_transfer_users,
+)
+
 package_group(
     name = "mosaic_users",
     includes = ["//jax:internal"],

jaxlib/jax.bzl

@@ -45,6 +45,7 @@ tf_cuda_tests_tags = _tf_cuda_tests_tags
 
 jax_internal_packages = []
 jax_extend_internal_users = []
+experimental_transfer_users = []
 mosaic_gpu_internal_users = []
 mosaic_internal_users = []
 pallas_gpu_internal_users = []