Move _src/custom_transpose.py into its own BUILD rule

Creating smaller build rules enforces better organized dependency graphs in the JAX project, helps pytype propagate annotations correctly, and leads to improved build and iteration times.

PiperOrigin-RevId: 762414305

Author: Jake VanderPlas

jax/BUILD

@@ -307,7 +307,6 @@ py_library_providing_imports_info(
         "_src/custom_derivatives.py",
         "_src/custom_partitioning.py",
         "_src/custom_partitioning_sharding_rule.py",
-        "_src/custom_transpose.py",
         "_src/debugging.py",
         "_src/dispatch.py",
         "_src/dlpack.py",
@@ -391,6 +390,7 @@ py_library_providing_imports_info(
         ":config",
         ":core",
         ":custom_api_util",
+        ":custom_transpose",
         ":deprecations",
         ":dtypes",
         ":effects",
@@ -595,6 +595,25 @@ pytype_strict_library(
     srcs = ["_src/custom_api_util.py"],
 )
 
+pytype_strict_library(
+    name = "custom_transpose",
+    srcs = ["_src/custom_transpose.py"],
+    deps = [
+        ":ad",
+        ":ad_util",
+        ":api_util",
+        ":core",
+        ":custom_api_util",
+        ":mlir",
+        ":partial_eval",
+        ":source_info_util",
+        ":traceback_util",
+        ":tree_util",
+        ":util",
+        ":xla",
+    ],
+)
+
 pytype_strict_library(
     name = "deprecations",
     srcs = ["_src/deprecations.py"],

tests/BUILD

@@ -61,6 +61,7 @@ jax_multiplatform_test(
     srcs = ["debug_info_test.py"],
     enable_configs = ["tpu_v3_x4"],
     deps = [
+        "//jax:custom_transpose",
         "//jax:experimental",
         "//jax:pallas",
         "//jax:pallas_gpu",