feature_request: support for tabulate/summary in the NNX API
Open
Priority: P1 - soon
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
Hey! I see that the Flax NNX APIs have visualization support via Penzai, but is there Module.tabulate or summary or something similar for nnx.Module objects? If not, is something already in the works?
Such functionalities are really helpful in quickly checking the number of parameters and distributions as well as the structure of the network and in debugging.
Contributor guide
Assessment
This issue has not been assessed yet.