"pure_callback" is not exported from module "jax"

[carlosgmartin]

Description

$ 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. 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:

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')

$ pyright --version
pyright 1.1.360

[jakevdp]

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.

[jakevdp]

Maybe this is the right fix:

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.

[jakevdp]

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 c1f65fc) 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.

[jakevdp]

Should be fixed in #20937