google / flax

Flax is a neural network library for JAX that is designed for flexibility.

Home Page:https://flax.readthedocs.io

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Improve documentation in flax.core

marcvanzee opened this issue · comments

Some functions in flax.core are not documented well and users who would like to use Flax for more advanced use cases have been complaining about this.

Some functions do not have docstring at all (e.g., pack), while others do have good docstrings, but the code itself is almost entirely undocumented, which makes is quite difficult to follow (e.g. ,vmap, scan, vjp). I think both should be improved.

#1992 improves some docstrings and fixes lint errors in a few flax.core files.

Closing this since we should probably use more specific issues instead.