search

poets-ai / elegy

A High Level API for Deep Learning in JAX
https://poets-ai.github.io/elegy/
MIT License
469 stars 32 forks source link

Using Elegy with Equinox? #235

Open nalzok opened 2 years ago

nalzok commented 2 years ago

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 commented 2 years ago

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.