Reduce test_print_wgmma_tiled_layout shape size.

PiperOrigin-RevId: 835230765

Author: IllogicalMoose

tests/pallas/mosaic_gpu_test.py

@@ -1203,7 +1203,7 @@ def kernel(x_ref, o_ref):
     self.assertEqual(output(), "It works!\n")
 
   def test_print_wgmma_tiled_layout(self):
-    shape = (128, 64)
+    shape = (64, 64)
     size = math.prod(shape)
 
     @functools.partial(