[jaxlib] Generate .pyi files for the _jax extension via nanobind's stubgen

jax/_src/compiler.py

@@ -637,11 +637,11 @@ def _share_fdo_profiles(
     backend: xc.Client,
     global_client: lib._jax.DistributedRuntimeClient,
     min_process_id
-) -> bytes | None:
+) -> bytes:
   sym_name = computation.operation.attributes['sym_name']
   module_name = ir.StringAttr(sym_name).value
   fdo_profile = compile_options.executable_build_options.fdo_profile
-  if fdo_profile is None or len(fdo_profile) == 0:
+  if len(fdo_profile) == 0:
     return fdo_profile
 
   compile_options.executable_build_options.fdo_profile = b""

jax/_src/tree_util.py

@@ -25,6 +25,7 @@
 
 from jax._src import traceback_util
 from jax._src.lib import pytree
+from jax._src.lib import version as jaxlib_version
 from jax._src.util import safe_zip, set_module
 from jax._src.util import unzip2
 
@@ -123,7 +124,9 @@ def treedef_tuple(treedefs: Iterable[PyTreeDef]) -> PyTreeDef:
   See Also:
     - :func:`jax.tree_util.treedef_children`
   """
-  return pytree.tuple(default_registry, list(treedefs))
+  if jaxlib_version < (0, 8, 0):
+    return pytree.tuple(default_registry, list(treedefs))  # type: ignore
+  return pytree.treedef_tuple(default_registry, list(treedefs))
 
 
 @export

jaxlib/BUILD

@@ -323,8 +323,17 @@ nanobind_pywrap_extension(
 nanobind_pywrap_extension(
     name = "_jax",
     srcs = ["jax.cc"],
+    additional_stubgen_deps = [
+        "//third_party/py/numpy",
+        "//jaxlib/mlir:ir",
+    ],
+    enable_stub_generation = True,
     pytype_deps = py_deps(["numpy"]),
     pytype_srcs = glob(["_jax/*.pyi"]),
+    stub_replacement_patterns = {
+        "jax.jaxlib._jax.Array$": "Array: Any",
+        "jax.jaxlib._jax.ArrayImpl$": "ArrayImpl: Any",
+    },
     visibility = jax_visibility("jaxlib/_jax"),
     deps = [
         ":config",