google / google/flax

Outdated and Broken Documentation: Module & Pytree Guide

Open
#5,100 3 comments 2 reactions 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

## 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**
```python
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:

```python
class Linear(nnx.Pytree): # AttributeError: module 'flax.nnx' has no attribute 'Pytree'
```

### 3. **Missing `nnx.List` Container**
```python
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 :)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.