[Question] Is there a preferred way to visualise a network written as an eqx.Module?

[jwtkeeble]

Hi All,

Is there a preferred way to visualize neural networks written within the equinox package? I understand that there's a pretty print for the PyTree representation. But is there a way (or a preferred method) for visualising an equinox module as a computational graph? For example, something like torchviz in PyTorch?

[patrick-kidger]

Not that I ever use myself. I'm sure something must exist though! So if you find something please report back :)

There are various pytree pretty-printers but it sounds like you're really after a visualisation of the computation graph, i.e. the jaxpr.

[jwtkeeble]

Thanks for the quick response, @patrick-kidger ! I'll have a look and see what I find and if I find anything useful, I'll attach it here.