[Mosaic GPU] Reduce test size to make it run on GPUs with less SMEM (e.g. RTX Pro 6000 Blackwell)

PiperOrigin-RevId: 840262900

Author: dimitar-asenov

tests/mosaic/gpu_test.py

@@ -5267,10 +5267,10 @@ def body(ctx, out, scratch):
     self.assertArraysEqual(kernel(), expected)
 
   @parameterized.parameters(
-      ((4, 64, 128), [[0], [1], [2]], (4, 64, 128), False),
-      ((4, 64, 128), [[0], [1, 2], [3]], (4, 4, 16, 128), False),
-      ((4, 8, 16, 128), [[0], [1], [2, 3], [4]], (4, 8, 2, 8, 128), False),
-      ((4, 64, 128), [[0, 1], [2], [3]], (2, 2, 64, 128), True),
+      ((4, 64, 64), [[0], [1], [2]], (4, 64, 64), False),
+      ((4, 64, 64), [[0], [1, 2], [3]], (4, 4, 16, 64), False),
+      ((4, 8, 16, 64), [[0], [1], [2, 3], [4]], (4, 8, 2, 8, 64), False),
+      ((4, 64, 64), [[0, 1], [2], [3]], (2, 2, 64, 64), True),
   )
   def test_memref_expand_shape(
       self, input_shape, reassociation, output_shape, has_transforms