Simplify handling of type stubs in nanobind extension rules.

Pass pytype_srcs as data to the pybind_extension rule.

PiperOrigin-RevId: 745238783

Author: hawkinsp

jaxlib/pywrap.bzl

@@ -73,17 +73,11 @@ def nanobind_pywrap_extension(
         name = name,
         srcs = [src_cc_name],
         deps = [":" + lib_name],
+        data = pytype_srcs,
         linkopts = linkopts,
         visibility = visibility,
         default_deps = [],
         common_lib_packages = [
             "jaxlib",
         ],
     )
-
-    # Create a py_library with the type stubs as data, on which wheel builds can depend.
-    native.py_library(
-        name = name + "_type_stubs",
-        data = pytype_srcs,
-        deps = pytype_deps,
-    )

jaxlib/tools/BUILD.bazel

@@ -67,7 +67,7 @@ py_binary(
         "//jaxlib:jaxlib_binaries",
         "//jaxlib:setup.py",
         "//jaxlib/xla:xla_client.py",
-        "//jaxlib/xla:xla_extension_type_stubs",
+        "//jaxlib/xla:xla_extension",
         "@xla//xla/ffi/api:api.h",
         "@xla//xla/ffi/api:c_api.h",
         "@xla//xla/ffi/api:ffi.h",

jaxlib/triton/BUILD

@@ -37,7 +37,6 @@ pytype_strict_library(
         [],
         [
             "//jaxlib/mlir/_mlir_libs:_triton_ext",
-            "//jaxlib/mlir/_mlir_libs:_triton_ext_type_stubs",
         ],
     ),
 )