nnx.pmap documentation seems to be for vmap
Open
Priority: P1 - soon
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
Hi,
The [source code documentation](https://github.com/google/flax/blob/c12ba53eb5d4e09913ffe575768d6ef20ce2698d/flax/nnx/transforms/iteration.py#L441-L520) for `nnx.pmap` appears to be instead referring to nnx.vmap.
Contributor guide
Assessment
This issue has not been assessed yet.