[Docs]: Updated NNX graph API Reference

[ibbyml]

Updated NNX graph API Reference

This PR expands and restructures the NNX graph API documentation to make the graph system easier to understand and navigate.

Changes:

• Reorganized the graph API reference into clearly scoped sections (graph state APIs, state utilities, traversal, context management, variable utilities, etc.).
• Added detailed docstrings and usage descriptions for many previously undocumented or minimally documented graph functions, including:
  • flatten / unflatten
  • state, update, pop, clone, pure
  • graph traversal utilities (iter_graph, recursive_map, find_duplicates)
  • functional helpers (call)
  • context management utilities (split_context, merge_context, update_context)
• Clarified the roles of GraphDef, GraphState, GraphFlatState, and PureState and how they relate to split / merge and functional graph usage.
• Introduced small documentation‑only typing wrappers (via typing.Annotated) for graph state aliases so they can be documented as named concepts.

Fixes: (#5161)

Checklist

• This change is discussed in a Github issue/discussion ([docs]: Improve docs for NNX API Reference #5161).
• The documentation and docstrings adhere to the documentation guidelines.

[ibbyml]

Hi. Happy to adjust or revert any of the structural or example changes if there’s a preference. I plan to work on the NNX's object API reference if that's OK.

Thanks!

[vfdev-5]

Thanks for working on the docs improvements @ibbyml !
This PR has large changes, I'll try to review it, but there is however one thing which may block it from landing. Rewriting docs into another format can be very opinionated. So, it would be preferable to maybe split the PR into clear small pieces (adding new docstrings or improving existing ones) such that everyone is OK with and leave opinionated pieces to another PR where we can find a middle-ground. What do you think?

For example, the point:

Reorganized the graph API reference into clearly scoped sections (graph state APIs, state utilities, traversal, context management, variable utilities, etc.).

Personally, I'm not finding it more clear than the previous docs. For example, here we see graph header and then State header, it is unclear whether it is a submodule or something else. If you want to emphasize a common property we may need rewrite that in a more clear way.

What can done is that we may inspire from Jax ref API: https://docs.jax.dev/en/latest/jax.numpy.html

[ibbyml]

Thanks for the review @vfdev-5. Definitely makes sense. I'll split these up then we can individually work on how the docs should be formatted. Afterwards, I'll apply the same formatting choices to the other parts of API Reference docs. I also agree that maybe we should take inspiration from the Jax docs.

Closing this PR and I'll open smaller ones without the opinionated changes.