Using Elegy with Equinox?

[nalzok]

I am under the impression that Elegy is not really an NN framework by itself. Instead, it wraps NN frameworks like Treex/Equinox/Flax/Haiku to automate the training loop and provide handy utilities like elegy.Model.summary. If that is the case, could you provide a quick example of using Elegy with Equinox?

[cgarciae]

Hey @nalzok, this version of Elegy relies heavily on Treex to wrap other frameworks, checkout FlaxModule and HaikuModule e.g.

https://github.com/cgarciae/treex/blob/c35d68b261144ed9f04babc4434419ff576d36ad/treex/nn/flax_module.py#L16

Using arbitrary frameworks will be much easier with #236. I expect to roll a dev version of Elegy soon, if you want to try it out I'd be happy to help you get started.