Vadim BERTRAND

Hi, It's a bit of both, here is the simplified version of my implementon (I did not check if that runs): ```python from functools import partial from jax.experimental.shard_map import shard_map...