New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
"pure_callback" is not exported from module "jax" #20925
Comments
Thanks - I really wish there were a way to export a name while also renaming it, without having to use 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 Python typing is such a mess. |
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. |
Actually, that won't work becuase Nevertheless, my snark about export conventions still stands. |
Should be fixed in #20937 |
Description
See this line. Unlike the other exports,
import pure_callback_api as pure_callback
doesn't follow theimport X as X
redundancy pattern that tells type checkers that this is a re-export, as described in PEP 484:System info (python version, jaxlib version, accelerator, etc.)
The text was updated successfully, but these errors were encountered: