MOME Container Implementation does not support non-array Genotypes

[Lookatator]

Hi,

It seems the current implementation of the MOME Container does not work if the Genotype is a pytree more general than an array. For example the following line fails to run for Neural Network Genotypes:

QDax/qdax/core/containers/mome_repertoire.py

Line 314 in 45c4c2a

pareto_front_genotypes=cell_genotype.squeeze(axis=0),

We have a working fix in the lab; which we will clean before submitting it in a PR :)

[limbryan]

I think this has been fixed and addressed now with #143 right @felixchalumeau? Where we change the explicit type in the repertoire to account for non-array genotypes and arbitrary pytrees (

QDax/qdax/core/containers/mapelites_repertoire.py

Line 390 in 773e646

lambda x: jnp.zeros(shape=(num_centroids,) + x.shape, dtype=x.dtype),

)