Split along layers

[leogao2]

Is it possible to split it such that layers are split along some dimension of the mesh too?

For example:

Mesh shape: x:16,y:32 Layout: layers: x, hidden: y

If I had 32 layers, for example, I'd like the result to have 2 layers on the first slice of x, 2 layers on the next slice, etc. Ideally, something like GPipe where the forward and backward passes are pipelined so that 15/16ths of the devices don't sit idle would be preferable, but even being able to do the split naïvely would be useful.