[docs]: Improve docs for NNX API Reference

[ibbyml]

Hi,

I’d like to improve the NNX API reference docs (not Linen). While working on a smaller docs issue (#5107), I noticed a few recurring problems across the NNX API reference: incomplete indexing/toctrees, inconsistent Sphinx cross-references vs code formatting, and missing/unclear docstrings for user-facing symbols. Some pages also look copied from older docs and don’t match the current NNX surface.

---

Example:

• flax.nnx.unflatten in the graph section has a kwarg outer_index_outer_ref=None, but its docstring refers to index_ref_cache (which doesn’t exist).

---

To keep review manageable, I’ll do this as a series of small PRs, going section by section in the API Reference page. To keep the scope manageable I'll focus on docs + docstrings only (no behavioral changes).

Is it ok if I send small PRs starting with flax.nnx.graph?

Thanks!

[vfdev-5]

Hi @ibbyml, thanks for pointing that out. Yes, keeping the docs updated during the development can be a bit complex. If you can pointing out all these inconsistencies and send small fixes, it would be awesome.

The team is also working on how to improve the merging process on the approved PRs, but meanwhile having the fixes on GitHub would be very helpful!

[ayulockedin]

Hi @ibbyml and @vfdev-5, this sounds like a great initiative! To help parallelize the effort, I can take ownership of auditing the flax.nnx.nn module (Linear, Conv, Attention, etc.).

I'll check for the same inconsistencies (missing args, broken cross-refs, and outdated Linen-style docstrings) and submit them as separate, small PRs. Does that sound good?