The current implementation supports only sharding of tensor axes that have size divisible by the mesh axis size.
Boian Petkantchin