[MGPU] Doc fix and shape size fix for broadcast in WGSplatFragLayout.

We enforce len(source_shape) <= len(target_shape)

PiperOrigin-RevId: 845281808

Author: golechwierowicz

jax/experimental/mosaic/gpu/fragmented_array.py

@@ -616,10 +616,14 @@ class WGSplatFragLayout:
   def can_broadcast_to(self, shape) -> bool:
     """Check that the shape can be broadcast.
 
-    Only dimensions of size 1 can be broadcast. All other dimensions
-    must be the same as the argument shape.
+    Source rank must be not larger than the target rank and all the trailing
+    dimensions of the target shape must equal the dimensions of the source shape
+    (or source's one must be equal to 1).
     """
-    return all(dim1 == dim2 or dim1 == 1 for dim1, dim2 in zip(self.shape[::-1], shape[::-1]))
+    return len(self.shape) <= len(shape) and all(
+        dim1 == dim2 or dim1 == 1
+        for dim1, dim2 in zip(self.shape[::-1], shape[::-1])
+    )
 
   def registers_element_type(self, t: ir.Type) -> ir.Type:
     return t