Add representative batch size

Otherwise we can't shard on the batch dimension. With this change everything works up to executing the solver

Author: fmassa

examples/example_ds3.py

@@ -825,7 +825,7 @@ def init_weights(
     ),
 )
 
-bs = 1
+bs = 8 * mesh.shape[0]
 seqlen = 1024
 dim = 4096