Expose _split_by_mesh_axis directly in pw_jax.

This change exposes the `_split_by_mesh_axis` function as `pw_jax.split_by_mesh_axis` instead of exposing the entire `jaxlib._pathways` module. This provides a more focused API.

PiperOrigin-RevId: 852364503

Author: lukebaumann

pathwaysutils/experimental/split_by_mesh_axis.py

@@ -161,7 +161,7 @@ def split_by_mesh_axis(
       for x in flat_arrays
   ]
 
-  flat_split_arrays = pw_jax.jaxlib_pathways._split_by_mesh_axis(  # pylint: disable=protected-access
+  flat_split_arrays = pw_jax.split_by_mesh_axis(
       arrays=flat_arrays,
       sharded_dim_idxs=sharded_dim_idxs,
       mesh_axis_sizes=mesh.axis_sizes,

pathwaysutils/jax/__init__.py

@@ -21,21 +21,21 @@
 import jax
 
 
-class _FakeJaxModule:
-  """A fake module that raises an ImportError when accessed.
+class _FakeJaxFunction:
+  """An object that raises an ImportError for __getattr__ and __call__.
 
-  This is used to provide a placeholder for JAX modules that are not available
-  in older versions of JAX, raising a helpful error message if they are
-  inadvertently used.
+  This is used to provide a placeholder for JAX functions that are not
+  available in older versions of JAX, raising a helpful error message if they
+  are inadvertently used.
   """
 
   def __init__(self, name, version):
     self.__name__ = name
     self.version = version
     self.error_message = (
-        f"Module {self.__name__} does not exist until JAX {self.version}. "
+        f"Function {self.__name__} does not exist until JAX {self.version}. "
         f"The current version of JAX is {jax.__version__}. "
-        "Using this modules results in this runtime error."
+        "Using this function results in this runtime error."
     )
 
   def __getattr__(self, name):
@@ -77,14 +77,20 @@ def register_backend_cache(cache: Any, name: str, util=util):  # pylint: disable
 
 try:
   # jax>=0.8.0
-  from jaxlib import _pathways as jaxlib_pathways  # pylint: disable=g-import-not-at-top
+  from jaxlib import _pathways  # pylint: disable=g-import-not-at-top
+
+  split_by_mesh_axis = _pathways._split_by_mesh_axis
+  del _pathways
 
 except ImportError:
   # jax<0.8.0
 
-  jaxlib_pathways = _FakeJaxModule("jax.jaxlib._pathways", "0.8.0")
+  split_by_mesh_axis = _FakeJaxFunction(
+      "jax.jaxlib._pathways._split_by_mesh_axis",
+      "0.8.0",
+  )
 
 
 del jax
 del Any
-del _FakeJaxModule
+del _FakeJaxFunction