feature_request: support for tabulate/summary in the NNX API

[AshishKumar4]

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.

[cgarciae]

Hey! Having tabulate API is definitely on the horizon as having a top-tier debugging experience is very important for NNX. Its not in the works but something that we would add sooner than later.

Thanks for bringing it up, we can track the issue here and hopefully we can spend some cycles on it soon.