Repository navigation
Outdated and Broken Documentation: Module & Pytree Guide #5100
Description
Activity
- Thanks for the report! Can you provide the version of Flax you have ?…On Tue, Nov 18, 2025, 18:22 Egor Konovalov ***@***.***> wrote: *wtfnukee* created an issue (google/flax#5100) <#5100> Description The Module & Pytree guide <https://flax.readthedocs.io/en/stable/guides/pytree.html> in the Flax documentation is severely outdated and non-functional. Approximately half of the code cells in this guide fail to execute, making it impossible for users to learn from the examples. Environment - Flax version: [latest stable as of the documentation] - JAX version: [current] - Python version: 3.12 Issues Found Running through the notebook reveals multiple breaking issues: 1. *First Example Fails Immediately* class Linear(nnx.Module): def __init__(self, din, dout, rngs: nnx.Rngs): self.din, self.dout = din, dout self.kernel = nnx.Param(rngs.normal((din, dout))) rngs = nnx.Rngs(0)weights = Linear(2, 3, rngs=rngs) *Error:* TypeError: RngStream.__call__() takes 1 positional argument but 2 were given 2. *Missing nnx.Pytree Class* Multiple examples reference nnx.Pytree which doesn't exist in the current API: class Linear(nnx.Pytree): # AttributeError: module 'flax.nnx' has no attribute 'Pytree' 3. *Missing nnx.List Container* self.layers = nnx.List([...]) # AttributeError: module 'flax.nnx' has no attribute 'List' 4. *Missing Utility Functions* - nnx.is_data() - doesn't exist - nnx.find_duplicates() - doesn't exist - Various other API mismatches 5. *API Inconsistencies Throughout* The guide references an API surface that appears to be from an older or planned version of NNX that doesn't match the current implementation. Impact This is a critical documentation issue because: 1. First impressions matter: Users trying to learn Flax NNX hit immediate failures 2. Wastes developer time: Hours spent debugging what turns out to be doc issues 3. Erodes trust: When core documentation doesn't work, it raises questions about library stability 4. Blocks adoption: Potential users will simply move to alternatives with working docs Expected Behavior Documentation examples should: - Execute without errors - Use current API patterns - Match the installed version of Flax - Include version compatibility notes if APIs changed Additional Context I understand that JAX/Flax is evolving rapidly and hasn't reached 1.0 yet. However, having non-functional core documentation creates a significant barrier to adoption. Even if the API is unstable, the docs should accurately reflect the current state. Would appreciate if this could be prioritized - happy to help test updated examples if needed. I'd love to write updated parts myself, but it's hard with incomplete docs like this :) — Reply to this email directly, view it on GitHub <#5100>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AASYOHYALHTNJJVAGJHACBD35NIVLAVCNFSM6AAAAACMPMDJPWVHI2DSMVQWIX3LMV43ASLTON2WKOZTGYZTSMRSGEYDIMI> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>
I was using default
0.10.7version in Colab, but updated to0.12.0. It did fix most of errors!
Now it only fails at Dataclasses part (AttributeError: Module flax.nnx has no attribute 'dataclass')Actually, Colab link shows the notebook from
mainand not the latest stable one. If you check the docs for the latest stable flax: https://flax.readthedocs.io/en/stable/guides/pytree.html, there is not that section with@nnx.dataclass.
Correct link for colab should be : https://colab.research.google.com/github/google/flax/blob/v0.12.0/docs_nnx/guides/pytree.ipynb#scrollTo=9ca77d65Thanks for pointing out this, we definitely need to figure out how to fix this.
Reacted by Egor Konovalov
Description
The Module & Pytree guide in the Flax documentation is severely outdated and non-functional. Approximately half of the code cells in this guide fail to execute, making it impossible for users to learn from the examples.
Environment
Issues Found
Running through the notebook reveals multiple breaking issues:
1. First Example Fails Immediately
Error:
2. Missing
nnx.PytreeClassMultiple examples reference
nnx.Pytreewhich doesn't exist in the current API:3. Missing
nnx.ListContainer4. Missing Utility Functions
nnx.is_data()- doesn't existnnx.find_duplicates()- doesn't exist5. API Inconsistencies Throughout
The guide references an API surface that appears to be from an older or planned version of NNX that doesn't match the current implementation.
Impact
This is a critical documentation issue because:
Expected Behavior
Documentation examples should:
Additional Context
I understand that JAX/Flax is evolving rapidly and hasn't reached 1.0 yet. However, having non-functional core documentation creates a significant barrier to adoption. Even if the API is unstable, the docs should accurately reflect the current state.
Would appreciate if this could be prioritized - happy to help test updated examples if needed. I'd love to write updated parts myself, but it's hard with incomplete docs like this :)