_sample_in_masked_pareto_front does not return random_key

[Lookatator]

The function _sample_in_masked_pareto_front does not return an updated random_key (contrary to all other functions in the project).

QDax/qdax/core/containers/mome_repertoire.py

Lines 49 to 61 in 50b997f

	@partial(jax.jit, static_argnames=("num_samples",))
	def _sample_in_masked_pareto_front(
	self,
	pareto_front_genotypes: Genotype,
	mask: jnp.ndarray,
	random_key: RNGKey,
	) -> Genotype:
	"""Sample num_samples elements in masked pareto front.
	Note: do not retrieve a random key because this function
	is to be vmapped. The public method that uses this function
	will return a random key

This can probably be solved by using the in_axis and out_axis arguments of jax.vmap