Repository: google/jax

Issue:
"pure_callback" is not exported from module "jax"
### Description

```shell
$ echo "import jax; jax.pure_callback" >> test.py
$ pyright test.py                                
/Users/carlos/Downloads/test.py
  /Users/carlos/Downloads/test.py:2:17 - error: "pure_callback" is not exported from module "jax" (reportPrivateImportUsage)
1 error, 0 warnings, 0 informations
```

See [this line](https://github.com/google/jax/blob/main/jax/__init__.py#L121). Unlike the other exports, `import pure_callback_api as pure_callback` doesn't follow the `import X as X` redundancy pattern that tells type checkers that this is a re-export, as described in [PEP 484](https://peps.python.org/pep-0484/):

> Modules and variables imported into the stub are not considered exported from the stub unless the import uses the `import ... as ...` form or the equivalent `from ... import ... as ...` form. (UPDATE: To clarify, the intention here is that only names imported using the form `X as X` will be exported, i.e. the name before and after as must be the same.)

### System info (python version, jaxlib version, accelerator, etc.)

```
jax:    0.4.26
jaxlib: 0.4.26
numpy:  1.26.4
python: 3.12.2 (main, Feb  6 2024, 20:19:44) [Clang 15.0.0 (clang-1500.1.0.2.5)]
jax.devices (1 total, 1 local): [CpuDevice(id=0)]
process_count: 1
platform: uname_result(system='Darwin', node='Carloss-MacBook-Pro-2.local', release='23.2.0', version='Darwin Kernel Version 23.2.0: Wed Nov 15 21:54:05 PST 2023; root:xnu-10002.61.3~2/RELEASE_ARM64_T6031', machine='arm64')
```

```shell
$ pyright --version
pyright 1.1.360
```


Discussion:
Thanks - I really wish there were a way to export a name while also renaming it, without having to use `__all__`. I think the only possibility here is to rename `pure_callback_api` to `pure_callback` in `jax._src.callback`, so that it can be imported by its exported name, which frankly is pretty silly.

Also, the context of this statement in PEP 484 is in a discussion of stub files, so it's not clear to me that this was ever intended to apply to `*.py` files. I've never found a good answer to that, but in any case `import X as X` within `py` files seems to be the export convention that at least some type checkers have landed on.

Python typing is such a mess.
Maybe this is the right fix:
```python
from jax._src.callback import pure_callback_api
jax._src.callback.pure_callback = pure_callback_api
del pure_callback_api
from jax._src.callback import pure_callback as pure_callback
del jax._src.callback.pure_callback
```
It certainly seems like the only way to export the API without renaming it at its source.
Actually, that won't work becuase `jax._src.callback.pure_callback` already exists. `pure_callback_api` is a shim around it that is basically identical except for the docstring (added in https://github.com/google/jax/commit/c1f65fc8b21bd8a238af2b7c9ff4c90196277be3) I'm not sure what the intent of introducing that was, but we could probably just delete `pure_callback_api` and import `pure_callback as pure_callback` directly.

Nevertheless, my snark about export conventions still stands.

Resolving patch:
diff --git a/jax/__init__.py b/jax/__init__.py
--- a/jax/__init__.py
+++ b/jax/__init__.py
@@ -118,7 +118,7 @@
 from jax._src.api import pmap as pmap
 from jax._src.xla_bridge import process_count as process_count
 from jax._src.xla_bridge import process_index as process_index
-from jax._src.callback import pure_callback_api as pure_callback
+from jax._src.callback import pure_callback as pure_callback
 from jax._src.ad_checkpoint import checkpoint_wrapper as remat
 from jax._src.api import ShapeDtypeStruct as ShapeDtypeStruct
 from jax._src.api import value_and_grad as value_and_grad
diff --git a/jax/_src/callback.py b/jax/_src/callback.py
--- a/jax/_src/callback.py
+++ b/jax/_src/callback.py
@@ -257,14 +257,40 @@ def pure_callback(
     vectorized: bool = False,
     **kwargs: Any,
 ):
-  """Calls a pure Python callback.
+  """Calls a pure Python callback. Works under :func:`jit`/:func:`~vmap`/etc.
 
   For more explanation, see `External Callbacks`_.
 
+  ``pure_callback`` enables calling a Python function in JIT-ed JAX functions.
+  The input ``callback`` will be passed NumPy arrays in place of JAX arrays and
+  should also return NumPy arrays. Execution takes place on CPU, like any
+  Python+NumPy function.
+
+  The callback is treated as functionally pure, meaning it has no side-effects
+  and its output value depends only on its argument values. As a consequence, it
+  is safe to be called multiple times (e.g. when transformed by :func:`~vmap` or
+  :func:`~pmap`), or not to be called at all when e.g. the output of a
+  `jit`-decorated function has no data dependence on its value. Pure callbacks
+  may also be reordered if data-dependence allows.
+
+  When `vmap`-ed the behavior will depend on the value of the
+  ``vectorized`` keyword argument. When ``vectorized`` is ``True``, the callback
+  is assumed to obey
+  ``jax.vmap(callback)(xs) == callback(xs) == jnp.stack([callback(x) for x in xs])``.
+  Therefore, the callback will be called directly on batched inputs (where the
+  batch axes are the leading dimensions). Additionally, the callbacks should
+  return outputs that have corresponding leading batch axes. If not vectorized
+  ``callback`` will be mapped sequentially across the batched axis.
+  For example, if ``callback = lambda x, y: np.matmul(x, y)``, then we are free
+  to set ``vectorized=True`` because the ``np.matmul`` function handles
+  arbitrary leading batch dimensions.
+
   Args:
     callback: function to execute on the host. The callback is assumed to be a pure
       function (i.e. one without side-effects): if an impure function is passed, it
-      may behave in unexpected ways, particularly under transformation.
+      may behave in unexpected ways, particularly under transformation. The callable
+      will be passed PyTrees of arrays as arguments, and should return a PyTree of
+      arrays that matches ``result_shape_dtypes``.
     result_shape_dtypes: pytree whose leaves have ``shape`` and ``dtype`` attributes,
       whose structure matches the expected output of the callback function at runtime.
       :class:`jax.ShapeDtypeStruct` is often used to define leaf values.
@@ -301,75 +327,6 @@ def pure_callback(
   return tree_util.tree_unflatten(out_tree, out_flat)
 
 
-def pure_callback_api(
-    callback: Callable[..., Any],
-    result_shape_dtypes: Any,
-    *args: Any,
-    sharding: SingleDeviceSharding | None = None,
-    vectorized: bool = False,
-    **kwargs: Any,
-):
-  """Applies a functionally pure Python callable. Works under :func:`jit`/:func:`~pmap`/etc.
-
-  ``pure_callback`` enables calling a Python function in JIT-ed JAX functions.
-  The input ``callback`` will be passed NumPy arrays in place of JAX arrays and
-  should also return NumPy arrays. Execution takes place on CPU, like any
-  Python+NumPy function.
-
-  The callback is treated as functionally pure, meaning it has no side-effects
-  and its output value depends only on its argument values. As a consequence, it
-  is safe to be called multiple times (e.g. when transformed by :func:`~vmap` or
-  :func:`~pmap`), or not to be called at all when e.g. the output of a
-  `jit`-decorated function has no data dependence on its value. Pure callbacks
-  may also be reordered if data-dependence allows.
-
-  When :func:`~pmap`-ed, the pure callback will be called several times (one on each
-  axis of the map). When `vmap`-ed the behavior will depend on the value of the
-  ``vectorized`` keyword argument. When ``vectorized`` is ``True``, the callback
-  is assumed to obey
-  ``jax.vmap(callback)(xs) == callback(xs) == jnp.stack([callback(x) for x in xs])``.
-  Therefore, the callback will be called directly on batched inputs (where the
-  batch axes are the leading dimensions). Additionally, the callbacks should
-  return outputs that have corresponding leading batch axes. If not vectorized
-  ``callback`` will be mapped sequentially across the batched axis.
-  For example, if ``callback = lambda x, y: np.matmul(x, y)``, then we are free
-  to set ``vectorized=True`` because the ``np.matmul`` function handles
-  arbitrary leading batch dimensions.
-
-  Args:
-    callback: A Python callable. The callable will be passed PyTrees of NumPy
-      arrays as arguments, and should return a PyTree of NumPy arrays that
-      matches ``result_shape_dtypes``.
-    result_shape_dtypes: A PyTree with leaves that are objects with ``shape``
-      and ``dtype`` attributes which represent to the shapes and dtypes of the
-      value of ``callback`` applied to ``args`` and ``kwargs``.
-    *args: The positional arguments to the callback. Must be PyTrees of JAX
-      types.
-    sharding: optional sharding that specifies the device from which the
-      callback should be invoked.
-    vectorized: A boolean that indicates whether or not ``callback`` is
-      vectorized, meaning it can handle arrays with additional leading
-      dimensions. If ``vectorized`` is `True`, when the callback is mapped
-      via `jax.vmap`, it will be called directly on inputs with leading batch
-      dimensions instead of executing ``callback`` on each mapped input
-      individually. The callback should also return outputs batched across the
-      leading axis. By default, ``vectorized`` is ``False``.
-    **kwargs: The keyword arguments to the callback. Must be PyTrees of JAX
-      types.
-
-  Returns:
-    The value of ``callback(*args, **kwargs)``.
-  """
-  return pure_callback(
-      callback,
-      result_shape_dtypes,
-      *args,
-      sharding=sharding,
-      vectorized=vectorized,
-      **kwargs,
-  )
-
-
 # IO Callback
 
 io_callback_p = core.Primitive("io_callback")


Test patch: