PEP484-compatible export of names in Flax
manifest opened this issue · comments
Now that Flax provides a marker file according PEP 561, mypy is able to find type annotations. The issue is that nothing is exported at the moment – mypy requires public API of a package being exported using either __all__ or import ... as ... syntax.
python3 -m venv env
source env/bin/activate
pip install -U pip mypy
git clone git@github.com:google/flax.git
pip install ./flax
python -m mypy --install-types --non-interactive --strict -c "from flax import linen as nn; nn.Dense(1)"
<string>:1: error: Module has no attribute "Dense"
Found 1 error in 1 file (checked 1 source file)Currently, projects that are using Flax need to set--no-implicit-reexport=False.
There are some references on that matter.
- 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.
- However, as an exception to the previous bullet, all objects imported into a stub using from ... import * are considered exported.
- MyPy docs
--no-implicit-reexportalways treated as enabled for stub files.
JAX recently started to export type annotations, it would be great to have similar changes to Flax if possible.
