pmap with varying degree of parallelization

[mattiasmar]

Hello,

Is is possible to use pmap with unique number of devices used for parallelization for each layer? (E.g. for some matmuls use many devices for parallelization and for others use just a few or even just one TPU/GPU.)
If yes, there an example available demonstrating that? Does JAX automatically handle the gather and scatter operations?