Rework Model Class about apax HOT 3 CLOSED

It looks like the only shared parameter is n_species in GMNN.__init__, so imho GMNN(Descriptor, MLP) looks better and is probably easier to maintain.
https://github.com/GM-NN/gmnn-jax/blob/ec7f81e6f82b9b934c9e510dcbb05f00d8a5ba62/gmnn_jax/model/gmnn.py#L25-L62

from apax.

What do you think, @PythonFZ ?

from apax.

This is going to be a bit more involved than I initially thought. haiku requires all Modules to be initialized within an hk.transform. That means a construction like this GMNN(Descriptor, MLP) is not possible.
I think there are two options around this:

1. ad abstraction on top: Instead of passing the modules directly we could pass around factory classes.
2. try and see if Flax offers a more elegant solution.

I will investigate these options eventually.

from apax.