Improve documentation in flax.core
marcvanzee opened this issue · comments
Marc van Zee commented
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.
Marc van Zee commented
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.
Marc van Zee commented
#1992 improves some docstrings and fixes lint errors in a few flax.core
files.
Marc van Zee commented
Closing this since we should probably use more specific issues instead.