Remove some deprecated BUILD aliases, close visibility of some others.

PiperOrigin-RevId: 842764995

Author: hawkinsp

BUILD.bazel

@@ -30,21 +30,43 @@ wheel_sources(
     data_srcs = ["//jax"],
     py_srcs = [
         "//jax",
-        "//jax:compilation_cache",
-        "//jax:experimental",
         "//jax/example_libraries:example_libraries",
-        "//jax:experimental_colocated_python",
-        "//jax:experimental_sparse",
-        "//jax:experimental_buffer_callback",
-        "//jax:experimental_serialize_executable",
-        "//jax:pallas_experimental_gpu_ops",
-        "//jax:pallas_fuser",
-        "//jax:pallas_gpu_ops",
-        "//jax:pallas_mosaic_gpu",
-        "//jax:pallas_tpu_ops",
-        "//jax:pallas_triton",
-        "//jax:source_mapper",
-        "//jax:sparse_test_util",
+        "//jax/example_libraries:optimizers",
+        "//jax/example_libraries:stax",
+        "//jax/experimental:buffer_callback",
+        "//jax/experimental:checkify",
+        "//jax/experimental:colocated_python",
+        "//jax/experimental:compilation_cache",
+        "//jax/experimental:compute_on",
+        "//jax/experimental:custom_dce",
+        "//jax/experimental:custom_partitioning",
+        "//jax/experimental:fused",
+        "//jax/experimental:hijax",
+        "//jax/experimental:jet",
+        "//jax/experimental:layout",
+        "//jax/experimental:mesh_utils",
+        "//jax/experimental:multihost_utils",
+        "//jax/experimental:ode",
+        "//jax/experimental:pallas_experimental_gpu_ops",
+        "//jax/experimental:pallas_fuser",
+        "//jax/experimental:pallas_gpu_ops",
+        "//jax/experimental:pallas_mosaic_gpu",
+        "//jax/experimental:pallas_tpu_ops",
+        "//jax/experimental:pallas_triton",
+        "//jax/experimental:pjit",
+        "//jax/experimental:profiler",
+        "//jax/experimental:rnn",
+        "//jax/experimental:scheduling_groups",
+        "//jax/experimental:serialize_executable",
+        "//jax/experimental:shard_alike",
+        "//jax/experimental:shard_map",
+        "//jax/experimental:source_mapper",
+        "//jax/experimental:sparse_test_util",
+        "//jax/experimental:sparse",
+        "//jax/experimental:topologies",
+        "//jax/experimental:transfer",
+        "//jax/experimental:xla_metadata",
+        "//jax/experimental",
         "//jax/_src:lax_reference",
         "//jax/_src:internal_export_back_compat_test_util",
         "//jax/_src:internal_export_back_compat_test_data",

jax/BUILD

@@ -20,6 +20,7 @@ load(
     "jax_extend_internal_users",
     "jax_extra_deps",
     "jax_internal_packages",
+    "jax_visibility",
     "py_deps",
     "py_library_providing_imports_info",
     "pytype_library",
@@ -259,7 +260,7 @@ pytype_strict_library(
 # TODO(dsuo): remove these aliases/targets.
 pytype_strict_library(
     name = "experimental",
-    visibility = ["//visibility:public"],
+    visibility = jax_visibility("experimental_deprecated_alias"),
     deps = [
         ":jax",
         "//jax/example_libraries:optimizers",
@@ -288,190 +289,44 @@ pytype_strict_library(
     ],
 )
 
-alias(
-    name = "experimental_buffer_callback",
-    actual = "//jax/experimental:buffer_callback",
-    visibility = ["//jax/experimental:buffer_callback_users"],
-)
-
-alias(
-    name = "experimental_colocated_python",
-    actual = "//jax/experimental:colocated_python",
-    visibility = ["//visibility:public"],
-)
-
-alias(
-    name = "experimental_compute_on",
-    actual = "//jax/experimental:compute_on",
-    visibility = ["//visibility:public"],
-)
-
-alias(
-    name = "compilation_cache",
-    actual = "//jax/experimental:compilation_cache",
-    visibility = ["//visibility:public"],
-)
-
-alias(
-    name = "jet",
-    actual = "//jax/experimental:jet",
-    visibility = ["//visibility:public"],
-)
-
 alias(
     name = "mesh_utils",
     actual = "//jax/experimental:mesh_utils",
-    visibility = ["//visibility:public"],
-)
-
-alias(
-    name = "experimental_mesh_utils",
-    actual = "//jax/experimental:mesh_utils",
-    visibility = ["//visibility:public"],
-)
-
-alias(
-    name = "mosaic",
-    actual = "//jax/experimental:mosaic",
-    visibility = ["//jax/experimental:mosaic_users"],
-)
-
-alias(
-    name = "mosaic_gpu",
-    actual = "//jax/experimental:mosaic_gpu",
-    visibility = ["//jax/experimental:mosaic_gpu_users"],
-)
-
-alias(
-    name = "experimental_multihost_utils",
-    actual = "//jax/experimental:multihost_utils",
-    visibility = ["//visibility:public"],
-)
-
-alias(
-    name = "ode",
-    actual = "//jax/experimental:ode",
-    visibility = ["//visibility:public"],
+    visibility = jax_visibility("mesh_utils_deprecated_alias"),
 )
 
 alias(
     name = "pallas",
     actual = "//jax/experimental:pallas",
-    visibility = ["//visibility:public"],
+    visibility = jax_visibility("pallas_deprecated_alias"),
 )
 
 alias(
     name = "pallas_fuser",
     actual = "//jax/experimental:pallas_fuser",
-    visibility = ["//jax/experimental:pallas_fuser_users"],
-)
-
-alias(
-    name = "pallas_gpu",
-    actual = "//jax/experimental:pallas_gpu",
-    visibility = ["//jax/experimental:pallas_gpu_users"],
-)
-
-alias(
-    name = "pallas_gpu_ops",
-    actual = "//jax/experimental:pallas_gpu_ops",
-    visibility = ["//jax/experimental:pallas_gpu_users"],
+    visibility = jax_visibility("pallas_fuser_deprecated_alias"),
 )
 
 alias(
     name = "pallas_mosaic_gpu",
     actual = "//jax/experimental:pallas_mosaic_gpu",
-    visibility = ["//jax/experimental:mosaic_gpu_users"],
+    visibility = jax_visibility("pallas_mosaic_gpu_deprecated_alias"),
 )
 
 alias(
     name = "pallas_tpu",
     actual = "//jax/experimental:pallas_tpu",
-    visibility = ["//visibility:public"],
-)
-
-alias(
-    name = "pallas_tpu_ops",
-    actual = "//jax/experimental:pallas_tpu_ops",
-    visibility = ["//visibility:public"],
-)
-
-alias(
-    name = "pallas_triton",
-    actual = "//jax/experimental:pallas_triton",
-    visibility = ["//jax/experimental:pallas_gpu_users"],
-)
-
-alias(
-    name = "pallas_experimental_gpu_ops",
-    actual = "//jax/experimental:pallas_experimental_gpu_ops",
-    visibility = ["//jax/experimental:mosaic_gpu_users"],
-)
-
-alias(
-    name = "experimental_profiler",
-    actual = "//jax/experimental:profiler",
-    visibility = ["//visibility:public"],
-)
-
-alias(
-    name = "experimental_pjit",
-    actual = "//jax/experimental:pjit",
-    visibility = ["//visibility:public"],
-)
-
-alias(
-    name = "rnn",
-    actual = "//jax/experimental:rnn",
-    visibility = ["//visibility:public"],
-)
-
-alias(
-    name = "experimental_serialize_executable",
-    actual = "//jax/experimental:serialize_executable",
-    visibility = ["//jax/experimental:serialize_executable_users"],
-)
-
-alias(
-    name = "source_mapper",
-    actual = "//jax/experimental:source_mapper",
-    visibility = ["//visibility:public"],
-)
-
-alias(
-    name = "experimental_sparse",
-    actual = "//jax/experimental:sparse",
-    visibility = ["//visibility:public"],
-)
-
-alias(
-    name = "sparse_test_util",
-    actual = "//jax/experimental:sparse_test_util",
-    visibility = [":internal"],
-)
-
-alias(
-    name = "experimental_topologies",
-    actual = "//jax/experimental:topologies",
-    visibility = ["//visibility:public"],
+    visibility = jax_visibility("pallas_tpu_deprecated_alias"),
 )
 
 alias(
     name = "experimental_transfer",
     actual = "//jax/experimental:transfer",
-    visibility = ["//jax/experimental:experimental_transfer_users"],
+    visibility = jax_visibility("experimental_transfer_deprecated_alias"),
 )
 
-# Aliases of example_library targets.
-# TODO(dsuo): remove these aliases.
 alias(
     name = "optimizers",
     actual = "//jax/example_libraries:optimizers",
-    visibility = ["//visibility:public"],
-)
-
-alias(
-    name = "stax",
-    actual = "//jax/example_libraries:stax",
-    visibility = ["//visibility:public"],
+    visibility = jax_visibility("optimizers_deprecated_alias"),
 )