Increase dimensions of dot operands to trigger the use of a custom ca…

…ll in XLA.

After openxla/xla#9319 landed in XLA, the dimensions
of the tested matrix multiplication were too small to trigger the use of a
custom call, thereby breaking the assertion below.

PiperOrigin-RevId: 615559772

tests/pgle_test.py

@@ -46,7 +46,7 @@ def f(x, y):
       z = x @ y
       return z @ y
 
-    shape = (8, 8)
+    shape = (16, 16)
     x = jnp.arange(math.prod(shape)).reshape(shape).astype(np.float32)
     y = x + 1
     f_lowered = f.lower(x, y)