Multiple devices meshes within the same jitted computation

[hr0nix]

Imagine the following hypothetical scenario. I have 16 devices. I also have some jitted computation, some parts of which I would like to shard as if my mesh was 8x2, but other parts should be sharded as if it was 4x4.

Is it true that currently there is no way to achieve what I need because jit assumes a single fixed device mesh? If so, what can I do to work around this limitation.

I can give some rationale for why I'd need to shard the computations in this fashion if needed, but basically it comes down to minimizing communications in a MoE-like model.

[hr0nix]

Ok, I've realized that I don't have to set up mesh globally, I can specify it with each sharding instead, which would allow me to pass different meshes for different computations.