flax-specific tools for diagramming a flax nn?

[grahamgower]

There's a nice list of diagramming tools for neural networks here, some of which work directly off models implemented for keras. Are there any such tools that do this for flax networks?

[marcvanzee]

Hi @grahamgower, you could convert you Flax module to TF using jax2tf and then use Tensorflow's diagramming tools in Tensorboard to visualize your model.

[grahamgower]

Thanks for the suggestion, but I find tensorboard's visuals to be more confusing than illuminating.

[marcvanzee]

Yes, they can be quite difficult to grasp, I agree. Another thing you could try is @cgarciae's recently introduced flax.linen.tabulate, which outputs Modules with various useful statistics in a tabular way, where you can select the depth of the submodule tree manually.

[grahamgower]

Ah, yes, tabulate is awesome and I'm already using it! And once jax-ml/jax#4085 is fixed, then it will be even better.

[marcvanzee]

Even better, when #2316 is in the submodules will be sorted by call order 😃 (it should be in one of these days).