Outdated and Broken Documentation: Module & Pytree Guide

[wtfnukee]

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

• 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 :)

[vfdev-5]

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: ***@***.***>

[wtfnukee]

I was using default 0.10.7 version in Colab, but updated to 0.12.0. It did fix most of errors!
Now it only fails at Dataclasses part (AttributeError: Module flax.nnx has no attribute 'dataclass')

[vfdev-5]

Actually, Colab link shows the notebook from main and 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=9ca77d65

Thanks for pointing out this, we definitely need to figure out how to fix this.