Lux.jl model visualizer

[NeroBlackstone]

Is there a tool like netron that can visualize the Lux model?
Or a convenient plot library that can plot the Lux model.

[avik-pal]

Once (if) lutzroeder/netron#1044, we will be able to export Lux models using reactant and load them directly in netron.

In the short term, updating https://github.com/rmsrosa/ChainPlots.jl to include lux as well would be very nice.

[wsmoses]

you can already use Reactant to export to tensorboard (e.g. I use the kernel profiling feature of with_profile to emit the log dir and generate EnzymeAD/Enzyme-JAX#256 (comment)). That's obviously the kernel view, but you can ask it to visualize the graph like in https://www.tensorflow.org/tensorboard/graphs .

Probably worth writing a reactant tutorial for profiling/graphs, x/ref EnzymeAD/Reactant.jl#581

[avik-pal]

Added an example in the docs on how to use with https://github.com/google-ai-edge/model-explorer